From 737200c7ef0f00b5d781fa8099e78d1af8257066 Mon Sep 17 00:00:00 2001
From: Jan Pochyla <jpochyla@gmail.com>
Date: Tue, 15 Nov 2016 13:46:00 +0100
Subject: [PATCH] debuglink: fix debuglink for V2

---
 trezorlib/debuglink.py | 2 ++
 trezorlib/transport.py | 6 ++++++
 2 files changed, 8 insertions(+)

diff --git a/trezorlib/debuglink.py b/trezorlib/debuglink.py
index 1f366e932c..4e55d1ecc4 100644
--- a/trezorlib/debuglink.py
+++ b/trezorlib/debuglink.py
@@ -15,11 +15,13 @@ def pprint(msg):
 class DebugLink(object):
     def __init__(self, transport, pin_func=pin_info, button_func=button_press):
         self.transport = transport
+        self.transport.session_begin()
 
         self.pin_func = pin_func
         self.button_func = button_func
 
     def close(self):
+        self.transport.session_end()
         self.transport.close()
 
     def _call(self, msg, nowait=False):
diff --git a/trezorlib/transport.py b/trezorlib/transport.py
index 201b35a973..b751e461bb 100644
--- a/trezorlib/transport.py
+++ b/trezorlib/transport.py
@@ -156,6 +156,9 @@ class TransportV1(Transport):
 
 class TransportV2(Transport):
     def write(self, msg):
+        if not self.session_id:
+            raise Exception('Missing session_id for v2 transport')
+
         data = bytearray(msg.SerializeToString())
 
         # Convert to unsigned in python2
@@ -185,6 +188,9 @@ class TransportV2(Transport):
             first = False
 
     def _read(self):
+        if not self.session_id:
+            raise Exception('Missing session_id for v2 transport')
+
         chunk = self._read_chunk()
         (session_id, msg_type, datalen, data) = self.parse_first(chunk)
         payloadlen = datalen + 4  # For the checksum