Python client for GEX

TinyFrame.py 12KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404
  1. class TinyFrame:
  2. def __init__(self, peer:int=1):
  3. self.write = None # the writer function should be attached here
  4. self.id_listeners = {}
  5. self.type_listeners = {}
  6. self.fallback_listener = None
  7. self.peer = peer # the peer bit
  8. # ----------------------------- FRAME FORMAT ---------------------------------
  9. # The format can be adjusted to fit your particular application needs
  10. # If the connection is reliable, you can disable the SOF byte and checksums.
  11. # That can save up to 9 bytes of overhead.
  12. # ,-----+-----+-----+------+------------+- - - -+-------------,
  13. # | SOF | ID | LEN | TYPE | HEAD_CKSUM | DATA | DATA_CKSUM |
  14. # | 0-1 | 1-4 | 1-4 | 1-4 | 0-4 | ... | 0-4 | <- size (bytes)
  15. # '-----+-----+-----+------+------------+- - - -+-------------'
  16. #
  17. # SOF ......... start of frame, usually 0x01 (optional, configurable)
  18. # ID ......... the frame ID (MSb is the peer bit)
  19. # LEN ......... number of data bytes in the frame
  20. # TYPE ........ message type (used to run Type Listeners, pick any values you like)
  21. # HEAD_CKSUM .. header checksum
  22. #
  23. # DATA ........ LEN bytes of data (can be 0, in which case DATA_CKSUM is omitted as well)
  24. # DATA_CKSUM .. checksum, implemented as XOR of all preceding bytes in the message
  25. # !!! BOTH SIDES MUST USE THE SAME SETTINGS !!!
  26. # Settings can be adjusted by setting the properties after init
  27. # Adjust sizes as desired (1,2,4)
  28. self.ID_BYTES = 2
  29. self.LEN_BYTES = 2
  30. self.TYPE_BYTES = 1
  31. # Checksum type
  32. # ('none', 'xor', 'crc16, 'crc32'
  33. self.CKSUM_TYPE = 'xor'
  34. # Use a SOF byte to mark the start of a frame
  35. self.USE_SOF_BYTE = True
  36. # Value of the SOF byte (if TF_USE_SOF_BYTE == 1)
  37. self.SOF_BYTE = 0x01
  38. self.next_frame_id = 0
  39. self.reset_parser()
  40. self._CKSUM_BYTES = None # will be updated on first compose / accept
  41. def reset_parser(self):
  42. """
  43. Reset the parser to its initial state
  44. """
  45. # parser state: SOF, ID, LEN, TYPE, HCK, PLD, PCK
  46. self.ps = 'SOF'
  47. # buffer for receiving bytes
  48. self.rbuf = None
  49. # expected number of bytes to receive
  50. self.rlen = 0
  51. # buffer for payload or checksum
  52. self.rpayload = None
  53. # received frame
  54. self.rf = TF_Msg()
  55. def _calc_cksum_bytes(self):
  56. """
  57. Get nbr of bytes needed for the checksum
  58. """
  59. if self.CKSUM_TYPE == 'none' or self.CKSUM_TYPE is None:
  60. return 0
  61. elif self.CKSUM_TYPE == 'xor':
  62. return 1
  63. elif self.CKSUM_TYPE == 'crc16':
  64. return 2
  65. elif self.CKSUM_TYPE == 'crc32':
  66. return 4
  67. else:
  68. raise Exception("Bad cksum type!")
  69. def _cksum(self, buffer) -> int:
  70. """
  71. Compute a checksum of the given buffer.
  72. """
  73. if self.CKSUM_TYPE == 'none' or self.CKSUM_TYPE is None:
  74. return 0
  75. elif self.CKSUM_TYPE == 'xor':
  76. acc = 0
  77. for b in buffer:
  78. acc ^= b
  79. return (~acc) & ((1<<(self._CKSUM_BYTES*8))-1)
  80. elif self.CKSUM_TYPE == 'crc16':
  81. # TODO implement crc16
  82. raise Exception("CRC16 not implemented!")
  83. elif self.CKSUM_TYPE == 'crc32':
  84. # TODO implement crc32
  85. raise Exception("CRC32 not implemented!")
  86. else:
  87. raise Exception("Bad cksum type!")
  88. def _gen_frame_id(self) -> int:
  89. """
  90. Get a new frame ID
  91. """
  92. frame_id = self.next_frame_id
  93. self.next_frame_id += 1
  94. if self.next_frame_id > ((1<<(8*self.ID_BYTES-1))-1):
  95. self.next_frame_id = 0
  96. if self.peer == 1:
  97. frame_id |= 1<<(8*self.ID_BYTES-1)
  98. return frame_id
  99. def _pack(self, num:int, bytes:int) -> bytes:
  100. """
  101. Pack a number for a TF field
  102. """
  103. return num.to_bytes(bytes, byteorder='big', signed=False)
  104. def _unpack(self, buf) -> int:
  105. """
  106. Unpack a number from a TF field
  107. """
  108. return int.from_bytes(buf, byteorder='big', signed=False)
  109. def query(self, type:int, listener, pld=None, id:int=None):
  110. """
  111. Send a query. Returns its ID
  112. """
  113. (id, buf) = self._compose(type=type, pld=pld, id=id)
  114. if listener is not None:
  115. self.add_id_listener(id, listener)
  116. self.write(buf)
  117. return id
  118. def send(self, type:int, pld=None, id:int=None):
  119. """
  120. Like query, but with no listener. Returns the ID
  121. """
  122. return self.query(type=type, pld=pld, id=id, listener=None)
  123. def _compose(self, type:int, pld=None, id:int=None) -> tuple:
  124. """
  125. Compose a frame.
  126. frame_id can be an ID of an existing session, None for a new session.
  127. """
  128. if self._CKSUM_BYTES is None:
  129. self._CKSUM_BYTES = self._calc_cksum_bytes()
  130. if pld is None:
  131. pld = bytearray()
  132. if id is None:
  133. id = self._gen_frame_id()
  134. buf = bytearray()
  135. if self.USE_SOF_BYTE:
  136. buf.extend(self._pack(self.SOF_BYTE, 1))
  137. buf.extend(self._pack(id, self.ID_BYTES))
  138. buf.extend(self._pack(len(pld), self.LEN_BYTES))
  139. buf.extend(self._pack(type, self.TYPE_BYTES))
  140. if self._CKSUM_BYTES > 0:
  141. buf.extend(self._pack(self._cksum(buf), self._CKSUM_BYTES))
  142. if len(pld) > 0:
  143. buf.extend(pld)
  144. buf.extend(self._pack(self._cksum(pld), self._CKSUM_BYTES))
  145. return (id, buf)
  146. def accept(self, bytes):
  147. """
  148. Parse bytes received on the serial port
  149. """
  150. for b in bytes:
  151. self.accept_byte(b)
  152. def accept_byte(self, b:int):
  153. """
  154. Handle a received byte
  155. """
  156. # TODO this seems ripe for rewrite to avoid repetitive code
  157. if self._CKSUM_BYTES is None:
  158. self._CKSUM_BYTES = self._calc_cksum_bytes()
  159. if self.ps == 'SOF':
  160. if self.USE_SOF_BYTE:
  161. if b != self.SOF_BYTE:
  162. return
  163. self.rpayload = bytearray()
  164. self.rpayload.append(b)
  165. self.ps = 'ID'
  166. self.rlen = self.ID_BYTES
  167. self.rbuf = bytearray()
  168. if self.USE_SOF_BYTE:
  169. return
  170. if self.ps == 'ID':
  171. self.rpayload.append(b)
  172. self.rbuf.append(b)
  173. if len(self.rbuf) == self.rlen:
  174. self.rf.id = self._unpack(self.rbuf)
  175. self.ps = 'LEN'
  176. self.rlen = self.LEN_BYTES
  177. self.rbuf = bytearray()
  178. return
  179. if self.ps == 'LEN':
  180. self.rpayload.append(b)
  181. self.rbuf.append(b)
  182. if len(self.rbuf) == self.rlen:
  183. self.rf.len = self._unpack(self.rbuf)
  184. self.ps = 'TYPE'
  185. self.rlen = self.TYPE_BYTES
  186. self.rbuf = bytearray()
  187. return
  188. if self.ps == 'TYPE':
  189. self.rpayload.append(b)
  190. self.rbuf.append(b)
  191. if len(self.rbuf) == self.rlen:
  192. self.rf.type = self._unpack(self.rbuf)
  193. if self._CKSUM_BYTES > 0:
  194. self.ps = 'HCK'
  195. self.rlen = self._CKSUM_BYTES
  196. self.rbuf = bytearray()
  197. else:
  198. self.ps = 'PLD'
  199. self.rlen = self.rf.len
  200. self.rbuf = bytearray()
  201. return
  202. if self.ps == 'HCK':
  203. self.rbuf.append(b)
  204. if len(self.rbuf) == self.rlen:
  205. hck = self._unpack(self.rbuf)
  206. actual = self._cksum(self.rpayload)
  207. if hck != actual:
  208. print("[TF] Header checksum mismatch")
  209. self.reset_parser()
  210. else:
  211. if self.rf.len == 0:
  212. self.handle_rx_frame()
  213. self.reset_parser()
  214. else:
  215. self.ps = 'PLD'
  216. self.rlen = self.rf.len
  217. self.rbuf = bytearray()
  218. self.rpayload = bytearray()
  219. return
  220. if self.ps == 'PLD':
  221. self.rpayload.append(b)
  222. self.rbuf.append(b)
  223. if len(self.rbuf) == self.rlen:
  224. self.rf.data = self.rpayload
  225. if self._CKSUM_BYTES > 0:
  226. self.ps = 'PCK'
  227. self.rlen = self._CKSUM_BYTES
  228. self.rbuf = bytearray()
  229. else:
  230. self.handle_rx_frame()
  231. self.reset_parser()
  232. return
  233. if self.ps == 'PCK':
  234. self.rbuf.append(b)
  235. if len(self.rbuf) == self.rlen:
  236. pck = self._unpack(self.rbuf)
  237. actual = self._cksum(self.rpayload)
  238. if pck != actual:
  239. print("[TF] Payload checksum mismatch (given %x, computed %x)" % (pck, actual))
  240. print(self.rpayload)
  241. self.reset_parser()
  242. else:
  243. self.handle_rx_frame()
  244. self.reset_parser()
  245. return
  246. def handle_rx_frame(self):
  247. """
  248. Process a received and verified frame by calling a listener.
  249. """
  250. frame = self.rf
  251. if frame.id in self.id_listeners and self.id_listeners[frame.id] is not None:
  252. lst = self.id_listeners[frame.id]
  253. rv = lst['fn'](self, frame)
  254. if rv == TF.CLOSE or rv is None:
  255. self.id_listeners[frame.id] = None
  256. return
  257. elif rv == TF.RENEW:
  258. lst.age = 0
  259. return
  260. elif rv == TF.STAY:
  261. return
  262. # TF.NEXT lets another handler process it
  263. if frame.type in self.type_listeners and self.type_listeners[frame.type] is not None:
  264. lst = self.type_listeners[frame.type]
  265. rv = lst['fn'](self, frame)
  266. if rv == TF.CLOSE:
  267. self.type_listeners[frame.type] = None
  268. return
  269. elif rv != TF.NEXT:
  270. return
  271. if self.fallback_listener is not None:
  272. lst = self.fallback_listener
  273. rv = lst['fn'](self, frame)
  274. if rv == TF.CLOSE:
  275. self.fallback_listener = None
  276. def remove_id_listener(self, id:int):
  277. """
  278. Remove a ID listener
  279. """
  280. self.id_listeners[id] = None
  281. def add_id_listener(self, id:int, lst, lifetime:float=None):
  282. """
  283. Add a ID listener that expires in "lifetime" seconds
  284. listener function takes two arguments:
  285. tinyframe instance and a msg object
  286. """
  287. self.id_listeners[id] = {
  288. 'fn': lst,
  289. 'lifetime': lifetime, # TODO implement timeouts
  290. 'age': 0,
  291. }
  292. def add_type_listener(self, type:int, lst):
  293. """
  294. Add a type listener
  295. listener function takes two arguments:
  296. tinyframe instance and a msg object
  297. """
  298. self.type_listeners[type] = {
  299. 'fn': lst,
  300. }
  301. def add_fallback_listener(self, lst):
  302. """
  303. Add a fallback listener
  304. listener function takes two arguments:
  305. tinyframe instance and a msg object
  306. """
  307. self.fallback_listener = {
  308. 'fn': lst,
  309. }
  310. class TF_Msg:
  311. """ A TF message object """
  312. def __init__(self):
  313. self.data = bytearray()
  314. self.len = 0
  315. self.type = 0
  316. self.id = 0
  317. def __str__(self):
  318. return f"ID {self.id:X}h, type {self.type:X}h, len {self.len:d}, body: {self.data}"
  319. class TF:
  320. """ Constants """
  321. STAY = 'STAY'
  322. RENEW = 'RENEW'
  323. CLOSE = 'CLOSE'
  324. NEXT = 'NEXT'