usbmux.py (7974B)
1 #!/usr/bin/python 2 # -*- coding: utf-8 -*- 3 # 4 # usbmux.py - usbmux client library for Python 5 # 6 # Copyright (C) 2009 Hector Martin "marcan" <hector@marcansoft.com> 7 # 8 # This program is free software; you can redistribute it and/or modify 9 # it under the terms of the GNU General Public License as published by 10 # the Free Software Foundation, either version 2 or version 3. 11 # 12 # This program is distributed in the hope that it will be useful, 13 # but WITHOUT ANY WARRANTY; without even the implied warranty of 14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 15 # GNU General Public License for more details. 16 # 17 # You should have received a copy of the GNU General Public License 18 # along with this program; if not, write to the Free Software 19 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA 20 21 import socket, struct, select, sys 22 23 try: 24 import plistlib 25 haveplist = True 26 except: 27 haveplist = False 28 29 class MuxError(Exception): 30 pass 31 32 class MuxVersionError(MuxError): 33 pass 34 35 class SafeStreamSocket: 36 def __init__(self, address, family): 37 self.sock = socket.socket(family, socket.SOCK_STREAM) 38 self.sock.connect(address) 39 def send(self, msg): 40 totalsent = 0 41 while totalsent < len(msg): 42 sent = self.sock.send(msg[totalsent:]) 43 if sent == 0: 44 raise MuxError("socket connection broken") 45 totalsent = totalsent + sent 46 def recv(self, size): 47 msg = '' 48 while len(msg) < size: 49 chunk = self.sock.recv(size-len(msg)) 50 if chunk == '': 51 raise MuxError("socket connection broken") 52 msg = msg + chunk 53 return msg 54 55 class MuxDevice(object): 56 def __init__(self, devid, usbprod, serial, location): 57 self.devid = devid 58 self.usbprod = usbprod 59 self.serial = serial 60 self.location = location 61 def __str__(self): 62 return "<MuxDevice: ID %d ProdID 0x%04x Serial '%s' Location 0x%x>"%(self.devid, self.usbprod, self.serial, self.location) 63 64 class BinaryProtocol(object): 65 TYPE_RESULT = 1 66 TYPE_CONNECT = 2 67 TYPE_LISTEN = 3 68 TYPE_DEVICE_ADD = 4 69 TYPE_DEVICE_REMOVE = 5 70 VERSION = 0 71 def __init__(self, socket): 72 self.socket = socket 73 self.connected = False 74 75 def _pack(self, req, payload): 76 if req == self.TYPE_CONNECT: 77 return struct.pack("IH", payload['DeviceID'], payload['PortNumber']) + "\x00\x00" 78 elif req == self.TYPE_LISTEN: 79 return "" 80 else: 81 raise ValueError("Invalid outgoing request type %d"%req) 82 83 def _unpack(self, resp, payload): 84 if resp == self.TYPE_RESULT: 85 return {'Number':struct.unpack("I", payload)[0]} 86 elif resp == self.TYPE_DEVICE_ADD: 87 devid, usbpid, serial, pad, location = struct.unpack("IH256sHI", payload) 88 serial = serial.split("\0")[0] 89 return {'DeviceID': devid, 'Properties': {'LocationID': location, 'SerialNumber': serial, 'ProductID': usbpid}} 90 elif resp == self.TYPE_DEVICE_REMOVE: 91 devid = struct.unpack("I", payload)[0] 92 return {'DeviceID': devid} 93 else: 94 raise MuxError("Invalid incoming request type %d"%req) 95 96 def sendpacket(self, req, tag, payload={}): 97 payload = self._pack(req, payload) 98 if self.connected: 99 raise MuxError("Mux is connected, cannot issue control packets") 100 length = 16 + len(payload) 101 data = struct.pack("IIII", length, self.VERSION, req, tag) + payload 102 self.socket.send(data) 103 def getpacket(self): 104 if self.connected: 105 raise MuxError("Mux is connected, cannot issue control packets") 106 dlen = self.socket.recv(4) 107 dlen = struct.unpack("I", dlen)[0] 108 body = self.socket.recv(dlen - 4) 109 version, resp, tag = struct.unpack("III",body[:0xc]) 110 if version != self.VERSION: 111 raise MuxVersionError("Version mismatch: expected %d, got %d"%(self.VERSION,version)) 112 payload = self._unpack(resp, body[0xc:]) 113 return (resp, tag, payload) 114 115 class PlistProtocol(BinaryProtocol): 116 TYPE_RESULT = "Result" 117 TYPE_CONNECT = "Connect" 118 TYPE_LISTEN = "Listen" 119 TYPE_DEVICE_ADD = "Attached" 120 TYPE_DEVICE_REMOVE = "Detached" #??? 121 TYPE_PLIST = 8 122 VERSION = 1 123 def __init__(self, socket): 124 if not haveplist: 125 raise Exception("You need the plistlib module") 126 BinaryProtocol.__init__(self, socket) 127 128 def _pack(self, req, payload): 129 return payload 130 131 def _unpack(self, resp, payload): 132 return payload 133 134 def sendpacket(self, req, tag, payload={}): 135 payload['ClientVersionString'] = 'usbmux.py by marcan' 136 if isinstance(req, int): 137 req = [self.TYPE_CONNECT, self.TYPE_LISTEN][req-2] 138 payload['MessageType'] = req 139 payload['ProgName'] = 'tcprelay' 140 BinaryProtocol.sendpacket(self, self.TYPE_PLIST, tag, plistlib.writePlistToString(payload)) 141 def getpacket(self): 142 resp, tag, payload = BinaryProtocol.getpacket(self) 143 if resp != self.TYPE_PLIST: 144 raise MuxError("Received non-plist type %d"%resp) 145 payload = plistlib.readPlistFromString(payload) 146 return payload['MessageType'], tag, payload 147 148 class MuxConnection(object): 149 def __init__(self, socketpath, protoclass): 150 self.socketpath = socketpath 151 if sys.platform in ['win32', 'cygwin']: 152 family = socket.AF_INET 153 address = ('127.0.0.1', 27015) 154 else: 155 family = socket.AF_UNIX 156 address = self.socketpath 157 self.socket = SafeStreamSocket(address, family) 158 self.proto = protoclass(self.socket) 159 self.pkttag = 1 160 self.devices = [] 161 162 def _getreply(self): 163 while True: 164 resp, tag, data = self.proto.getpacket() 165 if resp == self.proto.TYPE_RESULT: 166 return tag, data 167 else: 168 raise MuxError("Invalid packet type received: %d"%resp) 169 def _processpacket(self): 170 resp, tag, data = self.proto.getpacket() 171 if resp == self.proto.TYPE_DEVICE_ADD: 172 self.devices.append(MuxDevice(data['DeviceID'], data['Properties']['ProductID'], data['Properties']['SerialNumber'], data['Properties']['LocationID'])) 173 elif resp == self.proto.TYPE_DEVICE_REMOVE: 174 for dev in self.devices: 175 if dev.devid == data['DeviceID']: 176 self.devices.remove(dev) 177 elif resp == self.proto.TYPE_RESULT: 178 raise MuxError("Unexpected result: %d"%resp) 179 else: 180 raise MuxError("Invalid packet type received: %d"%resp) 181 def _exchange(self, req, payload={}): 182 mytag = self.pkttag 183 self.pkttag += 1 184 self.proto.sendpacket(req, mytag, payload) 185 recvtag, data = self._getreply() 186 if recvtag != mytag: 187 raise MuxError("Reply tag mismatch: expected %d, got %d"%(mytag, recvtag)) 188 return data['Number'] 189 190 def listen(self): 191 ret = self._exchange(self.proto.TYPE_LISTEN) 192 if ret != 0: 193 raise MuxError("Listen failed: error %d"%ret) 194 def process(self, timeout=None): 195 if self.proto.connected: 196 raise MuxError("Socket is connected, cannot process listener events") 197 rlo, wlo, xlo = select.select([self.socket.sock], [], [self.socket.sock], timeout) 198 if xlo: 199 self.socket.sock.close() 200 raise MuxError("Exception in listener socket") 201 if rlo: 202 self._processpacket() 203 def connect(self, device, port): 204 ret = self._exchange(self.proto.TYPE_CONNECT, {'DeviceID':device.devid, 'PortNumber':((port<<8) & 0xFF00) | (port>>8)}) 205 if ret != 0: 206 raise MuxError("Connect failed: error %d"%ret) 207 self.proto.connected = True 208 return self.socket.sock 209 def close(self): 210 self.socket.sock.close() 211 212 class USBMux(object): 213 def __init__(self, socketpath=None): 214 if socketpath is None: 215 if sys.platform == 'darwin': 216 socketpath = "/var/run/usbmuxd" 217 else: 218 socketpath = "/var/run/usbmuxd" 219 self.socketpath = socketpath 220 self.listener = MuxConnection(socketpath, BinaryProtocol) 221 try: 222 self.listener.listen() 223 self.version = 0 224 self.protoclass = BinaryProtocol 225 except MuxVersionError: 226 self.listener = MuxConnection(socketpath, PlistProtocol) 227 self.listener.listen() 228 self.protoclass = PlistProtocol 229 self.version = 1 230 self.devices = self.listener.devices 231 def process(self, timeout=None): 232 self.listener.process(timeout) 233 def connect(self, device, port): 234 connector = MuxConnection(self.socketpath, self.protoclass) 235 return connector.connect(device, port) 236 237 if __name__ == "__main__": 238 mux = USBMux() 239 print "Waiting for devices..." 240 if not mux.devices: 241 mux.process(0.1) 242 while True: 243 print "Devices:" 244 for dev in mux.devices: 245 print dev 246 mux.process()