WIP

trezorlib/transport.py

@@ -9,34 +9,12 @@ class ConnectionError(Exception):
 
 class Transport(object):
     def __init__(self, device, *args, **kwargs):
+        print("Transport constructor")
         self.device = device
+        self.session_id = 0
         self.session_depth = 0
         self._open()
 
-    def _open(self):
-        raise NotImplementedException("Not implemented")
-
-    def _close(self):
-        raise NotImplementedException("Not implemented")
-
-    def _write(self, msg, protobuf_msg):
-        raise NotImplementedException("Not implemented")
-
-    def _read(self):
-        raise NotImplementedException("Not implemented")
-
-    def _session_begin(self):
-        pass
-
-    def _session_end(self):
-        pass
-
-    def ready_to_read(self):
-        """
-        Returns True if there is data to be read from the transport.  Otherwise, False.
-        """
-        raise NotImplementedException("Not implemented")
-
     def session_begin(self):
         """
         Apply a lock to the device in order to preform synchronous multistep "conversations" with the device.  For example, before entering the transaction signing workflow, one begins a session.  After the transaction is complete, the session may be ended.
@@ -64,16 +42,14 @@ def write(self, msg):
         """
         Write mesage to tansport.  msg should be a member of a valid `protobuf class <https://developers.google.com/protocol-buffers/docs/pythontutorial>`_ with a SerializeToString() method.
         """
-        ser = msg.SerializeToString()
-        header = struct.pack(">HL", mapping.get_type(msg), len(ser))
-        self._write(b"##" + header + ser, msg)
+        raise NotImplementedException("Not implemented")
 
     def read(self):
         """
         If there is data available to be read from the transport, reads the data and tries to parse it as a protobuf message.  If the parsing succeeds, return a protobuf object.
         Otherwise, returns None.
         """
-        if not self.ready_to_read():
+        if not self._ready_to_read():
             return None
 
         data = self._read()
@@ -84,7 +60,7 @@ def read(self):
 
     def read_blocking(self):
         """
-        Same as read, except blocks untill data is available to be read.
+        Same as read, except blocks until data is available to be read.
         """
         while True:
             data = self._read()
@@ -94,39 +70,115 @@ def read_blocking(self):
         return self._parse_message(data)
 
     def _parse_message(self, data):
-        (msg_type, data) = data
+        (session_id, msg_type, data) = data
+
+        # Raise exception if we get the response with
+        # unexpected session ID
+        self._check_session_id(session_id)
+
         if msg_type == 'protobuf':
             return data
         else:
+            print mapping.get_class(msg_type)
             inst = mapping.get_class(msg_type)()
-            inst.ParseFromString(data)
+            print inst, data
+            inst.ParseFromString(bytes(data))
             return inst
 
-    def _read_headers(self, read_f):
-        # Try to read headers until some sane value are detected
-        is_ok = False
-        while not is_ok:
-
-            # Align cursor to the beginning of the header ("##")
-            c = read_f.read(1)
-            i = 0
-            while c != b"#":
-                i += 1
-                if i >= 64:
-                    # timeout
-                    raise Exception("Timed out while waiting for the magic character")
-                c = read_f.read(1)
-
-            if read_f.read(1) != b"#":
-                # Second character must be # to be valid header
-                raise Exception("Second magic character is broken")
-
-            # Now we're most likely on the beginning of the header
-            try:
-                headerlen = struct.calcsize(">HL")
-                (msg_type, datalen) = struct.unpack(">HL", read_f.read(headerlen))
-                break
-            except:
-                raise Exception("Cannot parse header length")
+    def _check_session_id(self, session_id):
+        if self.session_id == 0:
+            # Let the device set the session ID
+            self.session_id = session_id
+        elif session_id != self.session_id:
+            # Session ID has been already set, but it differs from response
+            raise Exception("Session ID mismatch. Have %d, got %d" % (self.session_id, session_id))
+
+    # Functions to be implemented in specific transports:
+    def _open(self):
+        raise NotImplementedException("Not implemented")
+
+    def _close(self):
+        raise NotImplementedException("Not implemented")
+
+    def _write_chunk(self, chunk):
+        raise NotImplementedException("Not implemented")
+
+    def _read_chunk(self):
+        raise NotImplementedException("Not implemented")
+
+    def _ready_to_read(self):
+        """
+        Returns True if there is data to be read from the transport.  Otherwise, False.
+        """
+        raise NotImplementedException("Not implemented")
+
+    def _session_begin(self):
+        pass
+
+    def _session_end(self):
+        pass
+
+class TransportV1(Transport):
+    def write(self, msg):
+        ser = msg.SerializeToString()
+        header = struct.pack(">HL", mapping.get_type(msg), len(ser))
+        data = bytearray(b"##" + header + ser)
+
+        while len(data):
+            # Report ID, data padded to 63 bytes
+            chunk = b'?' + data[:63] + b'\0' * (63 - len(data[:63]))
+            self._write_chunk(chunk)
+            data = data[63:]
+
+    def _read(self):
+        chunk = self._read_chunk()
+        (msg_type, datalen, data) = self.parse_first(chunk)
+
+        while len(data) < datalen:
+            chunk = self._read_chunk()
+            data.extend(self.parse_next(chunk))
+
+        # Strip padding zeros
+        data = data[:datalen]
+        return (0, msg_type, data)
+
+    def parse_first(self, chunk):
+        if chunk[:3] != b"?##":
+            raise Exception("Unexpected magic characters")
+
+        try:
+            headerlen = struct.calcsize(">HL")
+            (msg_type, datalen) = struct.unpack(">HL", chunk[3:3 + headerlen])
+        except:
+            raise Exception("Cannot parse header length")
+
+        data = chunk[3 + headerlen:]
+        return (msg_type, datalen, data)
+
+    def parse_next(self, chunk):
+        if chunk[0:1] != b"?":
+            raise Exception("Unexpected magic characters")
+
+        return chunk[1:]
+
+class TransportV2(Transport):
+    def write(self, msg):
+        ser = msg.SerializeToString()
+        raise NotImplemented()
+
+    def _read(self):
+        pass
+
+    def read_headers(self, read_f):
+        c = read_f.read(2)
+        if c != b"?!":
+            raise Exception("Unexpected magic characters")
+
+        try:
+            headerlen = struct.calcsize(">HL")
+            (session_id, msg_type, datalen) = struct.unpack(">LLL", read_f.read(headerlen))
+        except:
+            raise Exception("Cannot parse header length")
 
-        return (msg_type, datalen)
+        print datalen
+        return (0, msg_type, datalen)