git.sesse.net Git - pistorm/blob - a314/files_pi/picmd.py

   1 #!/usr/bin/python3
   2 # -*- coding: utf-8 -*-
   3
   4 # Copyright (c) 2018-2021 Niklas Ekström
   5
   6 import select
   7 import sys
   8 import socket
   9 import threading
  10 import time
  11 import os
  12 import struct
  13 import pty
  14 import signal
  15 import termios
  16 import fcntl
  17 import logging
  18 import json
  19
  20 logging.basicConfig(format = '%(levelname)s, %(asctime)s, %(name)s, line %(lineno)d: %(message)s')
  21 logger = logging.getLogger(__name__)
  22 logger.setLevel(logging.DEBUG)
  23
  24 FS_CFG_FILE = 'a314/files_pi/a314fs.conf'
  25 PICMD_CFG_FILE = 'a314/files_pi/picmd.conf'
  26
  27 volume_paths = {}
  28 search_path = ''
  29 env_vars = {}
  30 sgr_map = {}
  31
  32 def load_cfg():
  33     with open(FS_CFG_FILE, 'rt') as f:
  34         cfg = json.load(f)
  35         devs = cfg['devices']
  36         for _, dev in devs.items():
  37             volume_paths[dev['volume']] = dev['path']
  38
  39     global search_path
  40     search_path = os.getenv('PATH')
  41
  42     with open(PICMD_CFG_FILE, 'rt') as f:
  43         cfg = json.load(f)
  44
  45         if 'paths' in cfg:
  46             search_path = ':'.join(cfg['paths']) + ':' + search_path
  47             os.environ['PATH'] = search_path
  48
  49         if 'env_vars' in cfg:
  50             for key, val in cfg['env_vars'].items():
  51                 env_vars[key] = val
  52
  53         if 'sgr_map' in cfg:
  54             for key, val in cfg['sgr_map'].items():
  55                 sgr_map[key] = str(val)
  56
  57 load_cfg()
  58
  59 MSG_REGISTER_REQ        = 1
  60 MSG_REGISTER_RES        = 2
  61 MSG_DEREGISTER_REQ      = 3
  62 MSG_DEREGISTER_RES      = 4
  63 MSG_READ_MEM_REQ        = 5
  64 MSG_READ_MEM_RES        = 6
  65 MSG_WRITE_MEM_REQ       = 7
  66 MSG_WRITE_MEM_RES       = 8
  67 MSG_CONNECT             = 9
  68 MSG_CONNECT_RESPONSE    = 10
  69 MSG_DATA                = 11
  70 MSG_EOS                 = 12
  71 MSG_RESET               = 13
  72
  73 def wait_for_msg():
  74     header = b''
  75     while len(header) < 9:
  76         data = drv.recv(9 - len(header))
  77         if not data:
  78             logger.error('Connection to a314d was closed, terminating.')
  79             exit(-1)
  80         header += data
  81     (plen, stream_id, ptype) = struct.unpack('=IIB', header)
  82     payload = b''
  83     while len(payload) < plen:
  84         data = drv.recv(plen - len(payload))
  85         if not data:
  86             logger.error('Connection to a314d was closed, terminating.')
  87             exit(-1)
  88         payload += data
  89     return (stream_id, ptype, payload)
  90
  91 def send_register_req(name):
  92     m = struct.pack('=IIB', len(name), 0, MSG_REGISTER_REQ) + name
  93     drv.sendall(m)
  94
  95 def send_read_mem_req(address, length):
  96     m = struct.pack('=IIBII', 8, 0, MSG_READ_MEM_REQ, address, length)
  97     drv.sendall(m)
  98
  99 def read_mem(address, length):
 100     send_read_mem_req(address, length)
 101     stream_id, ptype, payload = wait_for_msg()
 102     if ptype != MSG_READ_MEM_RES:
 103         logger.error('Expected MSG_READ_MEM_RES but got %s. Shutting down.', ptype)
 104         exit(-1)
 105     return payload
 106
 107 def send_write_mem_req(address, data):
 108     m = struct.pack('=IIBI', 4 + len(data), 0, MSG_WRITE_MEM_REQ, address) + data
 109     drv.sendall(m)
 110
 111 def write_mem(address, data):
 112     send_write_mem_req(address, data)
 113     stream_id, ptype, payload = wait_for_msg()
 114     if ptype != MSG_WRITE_MEM_RES:
 115         logger.error('Expected MSG_WRITE_MEM_RES but got %s. Shutting down.', ptype)
 116         exit(-1)
 117
 118 def send_connect_response(stream_id, result):
 119     m = struct.pack('=IIBB', 1, stream_id, MSG_CONNECT_RESPONSE, result)
 120     drv.sendall(m)
 121
 122 def send_data(stream_id, data):
 123     m = struct.pack('=IIB', len(data), stream_id, MSG_DATA) + data
 124     drv.sendall(m)
 125
 126 def send_eos(stream_id):
 127     m = struct.pack('=IIB', 0, stream_id, MSG_EOS)
 128     drv.sendall(m)
 129
 130 def send_reset(stream_id):
 131     m = struct.pack('=IIB', 0, stream_id, MSG_RESET)
 132     drv.sendall(m)
 133
 134 sessions = {}
 135
 136 class PiCmdSession(object):
 137     def __init__(self, stream_id):
 138         self.stream_id = stream_id
 139         self.pid = 0
 140
 141         self.first_packet = True
 142         self.reset_after = None
 143
 144         self.rasp_was_esc = False
 145         self.rasp_in_cs = False
 146         self.rasp_holding = ''
 147
 148         self.amiga_in_cs = False
 149         self.amiga_holding = ''
 150
 151     def process_amiga_ansi(self, data):
 152         data = data.decode('latin-1')
 153         out = ''
 154         for c in data:
 155             if not self.amiga_in_cs:
 156                 if c == '\x9b':
 157                     self.amiga_in_cs = True
 158                     self.amiga_holding = '\x1b['
 159                 else:
 160                     out += c
 161             else: # self.amiga_in_cs
 162                 self.amiga_holding += c
 163                 if c >= chr(0x40) and c <= chr(0x7e):
 164                     if c == 'r':
 165                         # Window Bounds Report
 166                         # ESC[1;1;rows;cols r
 167                         rows, cols = map(int, self.amiga_holding[6:-2].split(';'))
 168                         winsize = struct.pack('HHHH', rows, cols, 0, 0)
 169                         fcntl.ioctl(self.fd, termios.TIOCSWINSZ, winsize)
 170                     elif c == '|':
 171                         # Input Event Report
 172                         # ESC[12;0;0;x;x;x;x;x|
 173                         # Window resized
 174                         send_data(self.stream_id, b'\x9b' + b'0 q')
 175                     else:
 176                         out += self.amiga_holding
 177                     self.amiga_holding = ''
 178                     self.amiga_in_cs = False
 179         if len(out) != 0:
 180             os.write(self.fd, out.encode('utf-8'))
 181
 182     def process_msg_data(self, data):
 183         if self.first_packet:
 184             if len(data) != 8:
 185                 send_reset(self.stream_id)
 186                 del sessions[self.stream_id]
 187             else:
 188                 address, length = struct.unpack('>II', data)
 189                 buf = read_mem(address, length)
 190
 191                 ind = 0
 192                 rows, cols = struct.unpack('>HH', buf[ind:ind+4])
 193                 ind += 4
 194
 195                 component_count = buf[ind]
 196                 ind += 1
 197
 198                 components = []
 199                 for _ in range(component_count):
 200                     n = buf[ind]
 201                     ind += 1
 202                     components.append(buf[ind:ind+n].decode('latin-1'))
 203                     ind += n
 204
 205                 arg_count = buf[ind]
 206                 ind += 1
 207
 208                 args = []
 209                 for _ in range(arg_count):
 210                     n = buf[ind]
 211                     ind += 1
 212                     args.append(buf[ind:ind+n].decode('latin-1'))
 213                     ind += n
 214
 215                 if arg_count == 0:
 216                     args.append('bash')
 217
 218                 self.pid, self.fd = pty.fork()
 219                 if self.pid == 0:
 220                     for key, val in env_vars.items():
 221                         os.putenv(key, val)
 222                     os.putenv('PATH', search_path)
 223                     os.putenv('TERM', 'ansi')
 224                     winsize = struct.pack('HHHH', rows, cols, 0, 0)
 225                     fcntl.ioctl(sys.stdin, termios.TIOCSWINSZ, winsize)
 226                     if component_count != 0 and components[0] in volume_paths:
 227                         path = volume_paths[components[0]]
 228                         os.chdir(os.path.join(path, *components[1:]))
 229                     else:
 230                         os.chdir(os.getenv('HOME', '/'))
 231                     os.execvp(args[0], args)
 232
 233                 self.first_packet = False
 234
 235         elif self.pid:
 236             self.process_amiga_ansi(data)
 237
 238     def close(self):
 239         if self.pid:
 240             os.kill(self.pid, signal.SIGTERM)
 241             self.pid = 0
 242             os.close(self.fd)
 243         del sessions[self.stream_id]
 244
 245     def process_rasp_ansi(self, text):
 246         text = text.decode('utf-8')
 247         out = ''
 248         for c in text:
 249             if not self.rasp_in_cs:
 250                 if not self.rasp_was_esc:
 251                     if c == '\x1b':
 252                         self.rasp_was_esc = True
 253                     else:
 254                         out += c
 255                 else: # self.rasp_was_esc
 256                     if c == '[':
 257                         self.rasp_was_esc = False
 258                         self.rasp_in_cs = True
 259                         self.rasp_holding = '\x1b['
 260                     elif c == '\x1b':
 261                         out += '\x1b'
 262                     else:
 263                         out += '\x1b'
 264                         out += c
 265                         self.rasp_was_esc = False
 266             else: # self.rasp_in_cs
 267                 self.rasp_holding += c
 268                 if c >= chr(0x40) and c <= chr(0x7e):
 269                     if c == 'm':
 270                         # Select Graphic Rendition
 271                         # ESC[30;37m
 272                         attrs = self.rasp_holding[2:-1].split(';')
 273                         attrs = [sgr_map[a] if a in sgr_map else a for a in attrs]
 274                         out += '\x1b[' + (';'.join(attrs)) + 'm'
 275                     else:
 276                         out += self.rasp_holding
 277                     self.rasp_holding = ''
 278                     self.rasp_in_cs = False
 279         return out.encode('latin-1', 'replace')
 280
 281     def handle_text(self):
 282         try:
 283             text = os.read(self.fd, 1024)
 284             text = self.process_rasp_ansi(text)
 285             while len(text) > 0:
 286                 take = min(len(text), 252)
 287                 send_data(self.stream_id, text[:take])
 288                 text = text[take:]
 289         except:
 290             #os.close(self.fd)
 291             os.kill(self.pid, signal.SIGTERM)
 292             self.pid = 0
 293             send_eos(self.stream_id)
 294             self.reset_after = time.time() + 10
 295
 296     def handle_timeout(self):
 297         if self.reset_after and self.reset_after < time.time():
 298             send_reset(self.stream_id)
 299             del sessions[self.stream_id]
 300
 301     def fileno(self):
 302         return self.fd
 303
 304 def process_drv_msg(stream_id, ptype, payload):
 305     if ptype == MSG_CONNECT:
 306         if payload == b'picmd':
 307             s = PiCmdSession(stream_id)
 308             sessions[stream_id] = s
 309             send_connect_response(stream_id, 0)
 310         else:
 311             send_connect_response(stream_id, 3)
 312     elif stream_id in sessions:
 313         s = sessions[stream_id]
 314
 315         if ptype == MSG_DATA:
 316             s.process_msg_data(payload)
 317         elif ptype == MSG_EOS:
 318             if s.pid:
 319                 send_eos(s.stream_id)
 320             s.close()
 321         elif ptype == MSG_RESET:
 322             s.close()
 323
 324 done = False
 325
 326 try:
 327     idx = sys.argv.index('-ondemand')
 328 except ValueError:
 329     idx = -1
 330
 331 if idx != -1:
 332     fd = int(sys.argv[idx + 1])
 333     drv = socket.socket(fileno=fd)
 334 else:
 335     drv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 336     drv.connect(('localhost', 7110))
 337     drv.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
 338
 339     send_register_req(b'picmd')
 340     _, _, payload = wait_for_msg()
 341     if payload[0] != 1:
 342         logger.error('Unable to register picmd with driver, shutting down')
 343         drv.close()
 344         done = True
 345
 346 rbuf = b''
 347
 348 if not done:
 349     logger.info('picmd server is running')
 350
 351 while not done:
 352     sel_fds = [drv] + [s for s in sessions.values() if s.pid]
 353     if idx == -1:
 354         sel_fds.append(sys.stdin)
 355     rfd, wfd, xfd = select.select(sel_fds, [], [], 5.0)
 356
 357     for fd in rfd:
 358         if fd == sys.stdin:
 359             line = sys.stdin.readline()
 360             if not line or line.startswith('quit'):
 361                 for s in sessions.values():
 362                     s.close()
 363                 drv.close()
 364                 done = True
 365         elif fd == drv:
 366             buf = drv.recv(1024)
 367             if not buf:
 368                 for s in sessions.values():
 369                     s.close()
 370                 drv.close()
 371                 done = True
 372             else:
 373                 rbuf += buf
 374                 while True:
 375                     if len(rbuf) < 9:
 376                         break
 377
 378                     (plen, stream_id, ptype) = struct.unpack('=IIB', rbuf[:9])
 379                     if len(rbuf) < 9 + plen:
 380                         break
 381
 382                     rbuf = rbuf[9:]
 383                     payload = rbuf[:plen]
 384                     rbuf = rbuf[plen:]
 385
 386                     process_drv_msg(stream_id, ptype, payload)
 387         else:
 388             fd.handle_text()
 389
 390     for s in sessions.values():
 391         s.handle_timeout()