1#!/usr/bin/env python3 2import socket 3import sys 4import select 5from select import EPOLLIN, EPOLLPRI, EPOLLERR 6import time 7import argparse 8 9TIMEOUT = 1.0 # seconds 10 11VERSION_HEADER = bytearray('MesaOverlayControlVersion', 'utf-8') 12DEVICE_NAME_HEADER = bytearray('DeviceName', 'utf-8') 13MESA_VERSION_HEADER = bytearray('MesaVersion', 'utf-8') 14 15DEFAULT_SERVER_ADDRESS = "\0mesa_overlay" 16 17class Connection: 18 def __init__(self, path): 19 # Create a Unix Domain socket and connect 20 sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) 21 try: 22 sock.connect(path) 23 except socket.error as msg: 24 print(msg) 25 sys.exit(1) 26 27 self.sock = sock 28 29 # initialize poll interface and register socket 30 epoll = select.epoll() 31 epoll.register(sock, EPOLLIN | EPOLLPRI | EPOLLERR) 32 self.epoll = epoll 33 34 def recv(self, timeout): 35 ''' 36 timeout as float in seconds 37 returns: 38 - None on error or disconnection 39 - bytes() (empty) on timeout 40 ''' 41 42 events = self.epoll.poll(timeout) 43 for ev in events: 44 (fd, event) = ev 45 if fd != self.sock.fileno(): 46 continue 47 48 # check for socket error 49 if event & EPOLLERR: 50 return None 51 52 # EPOLLIN or EPOLLPRI, just read the message 53 msg = self.sock.recv(4096) 54 55 # socket disconnected 56 if len(msg) == 0: 57 return None 58 59 return msg 60 61 return bytes() 62 63 def send(self, msg): 64 self.sock.send(msg) 65 66class MsgParser: 67 MSGBEGIN = bytes(':', 'utf-8')[0] 68 MSGEND = bytes(';', 'utf-8')[0] 69 MSGSEP = bytes('=', 'utf-8')[0] 70 71 def __init__(self, conn): 72 self.cmdpos = 0 73 self.parampos = 0 74 self.bufferpos = 0 75 self.reading_cmd = False 76 self.reading_param = False 77 self.buffer = None 78 self.cmd = bytearray(4096) 79 self.param = bytearray(4096) 80 81 self.conn = conn 82 83 def readCmd(self, ncmds, timeout=TIMEOUT): 84 ''' 85 returns: 86 - None on error or disconnection 87 - bytes() (empty) on timeout 88 ''' 89 90 parsed = [] 91 92 remaining = timeout 93 94 while remaining > 0 and ncmds > 0: 95 now = time.monotonic() 96 97 if self.buffer is None: 98 self.buffer = self.conn.recv(remaining) 99 self.bufferpos = 0 100 101 # disconnected or error 102 if self.buffer is None: 103 return None 104 105 for i in range(self.bufferpos, len(self.buffer)): 106 c = self.buffer[i] 107 self.bufferpos += 1 108 if c == self.MSGBEGIN: 109 self.cmdpos = 0 110 self.parampos = 0 111 self.reading_cmd = True 112 self.reading_param = False 113 elif c == self.MSGEND: 114 if not self.reading_cmd: 115 continue 116 self.reading_cmd = False 117 self.reading_param = False 118 119 cmd = self.cmd[0:self.cmdpos] 120 param = self.param[0:self.parampos] 121 self.reading_cmd = False 122 self.reading_param = False 123 124 parsed.append((cmd, param)) 125 ncmds -= 1 126 if ncmds == 0: 127 break 128 elif c == self.MSGSEP: 129 if self.reading_cmd: 130 self.reading_param = True 131 else: 132 if self.reading_param: 133 self.param[self.parampos] = c 134 self.parampos += 1 135 elif self.reading_cmd: 136 self.cmd[self.cmdpos] = c 137 self.cmdpos += 1 138 139 # if we read the entire buffer and didn't finish the command, 140 # throw it away 141 self.buffer = None 142 143 # check if we have time for another iteration 144 elapsed = time.monotonic() - now 145 remaining = max(0, remaining - elapsed) 146 147 # timeout 148 return parsed 149 150def control(args): 151 if args.socket: 152 address = '\0' + args.socket 153 else: 154 address = DEFAULT_SERVER_ADDRESS 155 156 conn = Connection(address) 157 msgparser = MsgParser(conn) 158 159 version = None 160 name = None 161 mesa_version = None 162 163 msgs = msgparser.readCmd(3) 164 165 for m in msgs: 166 cmd, param = m 167 if cmd == VERSION_HEADER: 168 version = int(param) 169 elif cmd == DEVICE_NAME_HEADER: 170 name = param.decode('utf-8') 171 elif cmd == MESA_VERSION_HEADER: 172 mesa_version = param.decode('utf-8') 173 174 if version != 1 or name is None or mesa_version is None: 175 print('ERROR: invalid protocol') 176 sys.exit(1) 177 178 if args.info: 179 print(f"Protocol Version: {version}") 180 print(f"Device Name: {name}") 181 print(f"Mesa Version: {mesa_version}") 182 183 if args.cmd == 'start-capture': 184 conn.send(bytearray(':capture=1;', 'utf-8')) 185 elif args.cmd == 'stop-capture': 186 conn.send(bytearray(':capture=0;', 'utf-8')) 187 188if __name__ == '__main__': 189 parser = argparse.ArgumentParser(description='MESA_overlay control client') 190 parser.add_argument('--info', action='store_true', help='Print info from socket') 191 parser.add_argument('--socket', '-s', type=str, help='Path to socket') 192 193 commands = parser.add_subparsers(help='commands to run', dest='cmd') 194 commands.add_parser('start-capture') 195 commands.add_parser('stop-capture') 196 197 args = parser.parse_args() 198 199 control(args) 200