vt_decor.py

   1 import base64
   2 import inspect
   3 import ipaddress
   4 import re
   5 import socket
   6 import textwrap
   7 import time
   8 import weakref
   9
  10 from .. import chunks, messages, roost, util, zcode
  11
  12 ## Config
  13 # I'm not sure why I can't just use the standard form of
  14 # _enabled = util.Configurable(...) here, but this works.
  15
  16 SETTINGS = {
  17     'enabled': True,
  18     'min_width': 100,
  19     'time': 'receiveTime',
  20     'show_seconds': False,
  21     'show_addr': False,
  22     'show_signatures': False,
  23 }
  24
  25 HOSTNAMES = {}
  26
  27 def _gen_setter(key):
  28     def setter(_context, value):
  29         SETTINGS[key] = value
  30     return setter
  31
  32 util.Configurable(
  33     'user.vt_decor.enabled',
  34     True,
  35     action=_gen_setter('enabled'),
  36     oneof=('True', 'False'),
  37     validate=lambda val: isinstance(val, bool),
  38     coerce=util.coerce_bool,
  39     )
  40
  41 util.Configurable(
  42     'user.vt_decor.time',
  43     'receiveTime',
  44     action=_gen_setter('time'),
  45     oneof=('time', 'receiveTime'),
  46     )
  47
  48 util.Configurable(
  49     'user.vt_decor.show_seconds',
  50     False,
  51     action=_gen_setter('show_seconds'),
  52     oneof=('True', 'False'),
  53     validate=lambda val: isinstance(val, bool),
  54     coerce=util.coerce_bool,
  55     )
  56
  57 util.Configurable(
  58     'user.vt_decor.show_addr',
  59     False,
  60     action=_gen_setter('show_addr'),
  61     oneof=('True', 'False'),
  62     validate=lambda val: isinstance(val, bool),
  63     coerce=util.coerce_bool,
  64     )
  65
  66 util.Configurable(
  67     'user.vt_decor.show_signatures',
  68     False,
  69     action=_gen_setter('show_signatures'),
  70     oneof=('True', 'False'),
  71     validate=lambda val: isinstance(val, bool),
  72     coerce=util.coerce_bool,
  73     )
  74
  75
  76 ## Utils
  77
  78 UI = lambda: None
  79
  80 def _get_width():
  81     global UI
  82     try:
  83         return UI().maxx
  84     except AttributeError:
  85         pass
  86
  87     frame = inspect.currentframe()
  88     try:
  89         while True:
  90             frame = frame.f_back
  91             try:
  92                 ui = frame.f_locals['self'].ui
  93                 UI = weakref.ref(ui)
  94                 return ui.maxx
  95             except (KeyError, AttributeError):
  96                 pass
  97     finally:
  98         del frame
  99
 100
 101 ## The fun part
 102
 103 MSG_NO_WRAP_PATTERNS = [re.compile(x, re.MULTILINE) for x in (
 104     r'^(?: |>)',
 105     r'\t',
 106     r'[^.?!,]  ',
 107 )]
 108
 109 MSG_LIST_WRAP_PATTERN = re.compile(r'(^\s*[-*]\s+)', re.MULTILINE)
 110
 111 REALM_MAP = {
 112     'ANDREW.CMU.EDU': 'AN',
 113     'CS.CMU.EDU': 'CS',
 114     'IASTATE.EDU': 'IA',
 115 }
 116
 117 class RoostVTDecor(roost.RoostMessage.Decor):
 118
 119
 120     @classmethod
 121     def decorate(cls, msg, decoration):
 122         try:
 123             if SETTINGS['enabled'] and _get_width() >= SETTINGS['min_width']:
 124                 return cls.vt_decorate(msg, decoration)
 125         except Exception:
 126             pass
 127         return super().decorate(msg, decoration)
 128
 129     @staticmethod
 130     def rewrap_p(msg):
 131         force_wrap_tuples = {
 132             ('moira', 'incremental'),
 133             ('scripts', 'nagios.multivalue-key.mysql-s'),
 134             ('scripts', 'nagios.unique-key.mysql-s'),
 135         }
 136
 137         msg_tuples = {
 138             (msg.data['classKey'], '*'),
 139             (msg.data['classKey'], msg.data['instanceKey']),
 140         }
 141
 142         return (bool(force_wrap_tuples & msg_tuples) or
 143                 not any(x.search(msg.body) for x in MSG_NO_WRAP_PATTERNS))
 144
 145     @staticmethod
 146     def list_rewrap_p(msg):
 147         return len(MSG_LIST_WRAP_PATTERN.findall(msg.body)) >= 2
 148
 149     @classmethod
 150     def partial_rewrap(cls, body, width, indent):
 151         try:
 152             ret = ''
 153             bindent = ''
 154             for m in MSG_LIST_WRAP_PATTERN.split(body):
 155                 if not m:
 156                     continue
 157                 if MSG_LIST_WRAP_PATTERN.match(m):
 158                     # Bullet
 159                     if ret:
 160                         ret += indent
 161                     ret += m.lstrip('\n')
 162                     bindent = ' ' * util.glyphwidth(m)
 163                 else:
 164                     # Content
 165                     # collapse whitespace and wrap.
 166                     ret += textwrap.fill(re.sub(r'\s+', ' ', m), width,
 167                                          initial_indent=indent+bindent,
 168                                          subsequent_indent=indent+bindent,
 169                                          break_long_words=False,
 170                                          break_on_hyphens=False).lstrip()
 171                     ret += '\n'
 172
 173             return ret.rstrip()
 174         except Exception as e:
 175             return str(e)
 176
 177
 178     @classmethod
 179     def vt_decorate(cls, msg, decoration):
 180         width = _get_width()
 181
 182         tags = cls.decotags(decoration)
 183
 184         realm = msg.backend.realm
 185         sender = msg.data['sender']
 186         recipient = msg.data['recipient']
 187         if msg.data['isPersonal'] and msg.data['isOutgoing']:
 188             sender = f'➤{recipient}'
 189         if sender.endswith(realm):
 190             sender = sender[:sender.index('@')]
 191         klass = msg.data['class'] if msg.data['classKey'] != 'message' else ''
 192         inst = msg.data['instance'] or "''"
 193         if klass and inst.lower() == 'personal':
 194             dest = klass
 195         else:
 196             dest = f'{klass}[{inst}]'
 197         auth = '+' if msg.data['auth'] else '-'
 198         if not msg.data['auth'] and  msg.data['opcode'] == 'mattermost':
 199             auth = '¤'
 200         t = time.strftime(
 201             '%H:%M:%S' if SETTINGS['show_seconds'] else '%H:%M',
 202             time.localtime(msg.data[SETTINGS['time']] / 1000))
 203
 204         if recipient.startswith('@'):
 205             mrealm = REALM_MAP.get(recipient[1:], '??')
 206             dest = f'{mrealm} {dest}'
 207
 208         dest_width = 18 - (util.glyphwidth(dest) - len(dest))
 209
 210         prefix = f'{sender:10.10} {t} {dest:{dest_width}.{dest_width}} {auth} '
 211         indent = ' ' * util.glyphwidth(prefix)
 212
 213         body = msg.body.rstrip()
 214         if cls.list_rewrap_p(msg):
 215             body = cls.partial_rewrap(body, width, indent)
 216         elif cls.rewrap_p(msg):
 217             body = textwrap.fill(body, width,
 218                                  initial_indent=indent,
 219                                  subsequent_indent=indent,
 220                                  break_long_words=False,
 221                                  break_on_hyphens=False).lstrip()
 222         else:
 223             body = textwrap.indent(body, indent).lstrip()
 224
 225         # body = f'{list(decoration.keys())}\n'
 226         if msg.backend.format_body == 'format':
 227             cbody = zcode.tag(body, frozenset(tags))
 228         elif msg.backend.format_body == 'clear':
 229             cbody = chunks.Chunk([(tags, '')])
 230         else:
 231             if msg.backend.format_body == 'strip':
 232                 body = zcode.strip(body)
 233             cbody = chunks.Chunk([(tags, body)])
 234
 235
 236         chunk = chunks.Chunk([(tags, f'{prefix}')]) + cbody
 237
 238         if SETTINGS['show_signatures'] and msg.backend.format_zsig != 'clear':
 239             zsig = '\n' + textwrap.fill(
 240                 msg.data['signature'], width,
 241                 initial_indent='                  -- ',
 242                 subsequent_indent='                     ',
 243                 break_long_words=False,
 244                 break_on_hyphens=False)
 245
 246             if msg.backend.format_zsig == 'format':
 247                 chunk += zcode.tag(zsig, frozenset(tags))
 248             elif msg.backend.format_zsig == 'strip':
 249                 chunk.append((tags, zcode.strip(zsig)))
 250             else:
 251                 chunk.append((tags, zsig))
 252
 253         if SETTINGS['show_addr']:
 254             addr = str(ipaddress.ip_address(base64.b64decode(msg.data["uid"])[:4]))
 255             # Hangs snipe while resolving; need to make async somehow.
 256             #if addr not in HOSTNAMES:
 257             #    try:
 258             #        HOSTNAMES[addr] = socket.gethostbyaddr(addr)[0].upper()
 259             #    except Exception:
 260             #        HOSTNAMES[addr] = addr
 261             #addr = HOSTNAMES[addr]
 262             chunk.append((tags | {'right'}, f'## {addr}'))
 263         else:
 264             chunk.append((tags | {'right'}, ''))
 265
 266         return chunk