From 2813522b05cef4e0e545a101f8b3559a3183b45b Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Mon, 15 Apr 2019 16:15:28 +0200
Subject: [PATCH] tools: refresh mem_* debug tools, blackify

---
 tools/mem_flashblock.py | 40 +++++++++++++++++++++-------------------
 tools/mem_read.py       | 35 ++++++++++++++++-------------------
 tools/mem_write.py      | 33 +++++++++++++++------------------
 3 files changed, 52 insertions(+), 56 deletions(-)

diff --git a/tools/mem_flashblock.py b/tools/mem_flashblock.py
index 1df65dead..cde9d6313 100755
--- a/tools/mem_flashblock.py
+++ b/tools/mem_flashblock.py
@@ -1,47 +1,49 @@
 #!/usr/bin/env python3
 from trezorlib.debuglink import DebugLink
-from trezorlib.client import TrezorClient
 from trezorlib.transport import enumerate_devices
 import sys
 
+# fmt: off
 sectoraddrs = [0x8000000, 0x8004000, 0x8008000, 0x800c000,
                0x8010000, 0x8020000, 0x8040000, 0x8060000,
                0x8080000, 0x80a0000, 0x80c0000, 0x80f0000]
 sectorlens = [0x4000, 0x4000, 0x4000, 0x4000,
               0x8000, 0x10000, 0x10000, 0x10000,
               0x10000, 0x10000, 0x10000, 0x10000]
+# fmt: on
 
 
-def main():
-    # List all debuggable TREZORs
-    devices = [device for device in enumerate_devices() if hasattr(device, 'find_debug')]
-
-    # Check whether we found any
-    if len(devices) == 0:
-        print('No TREZOR found')
-        return
+def find_debug():
+    for device in enumerate_devices():
+        try:
+            debug_transport = device.find_debug()
+            debug = DebugLink(debug_transport, auto_interact=False)
+            debug.open()
+            return debug
+        except Exception:
+            continue
+    else:
+        print("No suitable Trezor device found")
+        sys.exit(1)
 
-    # Use first connected device
-    transport = devices[0]
-    debug_transport = devices[0].find_debug()
 
-    # Creates object for manipulating TREZOR
-    client = TrezorClient(transport)
-    debug = DebugLink(debug_transport)
+def main():
+    debug = find_debug()
 
     sector = int(sys.argv[1])
     f = open(sys.argv[2], "rb")
     content = f.read(sectorlens[sector])
-    if (len(content) != sectorlens[sector]):
+    if len(content) != sectorlens[sector]:
         print("Not enough bytes in file")
         return
 
     debug.flash_erase(sector)
     step = 0x400
     for offset in range(0, sectorlens[sector], step):
-        debug.memory_write(sectoraddrs[sector] + offset, content[offset:offset + step], flash=True)
-    client.close()
+        debug.memory_write(
+            sectoraddrs[sector] + offset, content[offset : offset + step], flash=True
+        )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()
diff --git a/tools/mem_read.py b/tools/mem_read.py
index c6a4a3a16..844118871 100755
--- a/tools/mem_read.py
+++ b/tools/mem_read.py
@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 from trezorlib.debuglink import DebugLink
-from trezorlib.client import TrezorClient
 from trezorlib.transport import enumerate_devices
 import sys
 
@@ -13,28 +12,28 @@ import sys
 # be running a firmware that was built with debug link enabled
 
 
-def main():
-    # List all debuggable TREZORs
-    devices = [device for device in enumerate_devices() if hasattr(device, 'find_debug')]
-
-    # Check whether we found any
-    if len(devices) == 0:
-        print('No TREZOR found')
-        return
+def find_debug():
+    for device in enumerate_devices():
+        try:
+            debug_transport = device.find_debug()
+            debug = DebugLink(debug_transport, auto_interact=False)
+            debug.open()
+            return debug
+        except Exception:
+            continue
+    else:
+        print("No suitable Trezor device found")
+        sys.exit(1)
 
-    # Use first connected device
-    transport = devices[0]
-    debug_transport = devices[0].find_debug()
 
-    # Creates object for manipulating TREZOR
-    client = TrezorClient(transport)
-    debug = DebugLink(debug_transport)
+def main():
+    debug = find_debug()
 
     arg1 = int(sys.argv[1], 16)
     arg2 = int(sys.argv[2], 16)
     step = 0x400 if arg2 >= 0x400 else arg2
 
-    f = open('memory.dat', 'wb')
+    f = open("memory.dat", "wb")
 
     for addr in range(arg1, arg1 + arg2, step):
         mem = debug.memory_read(addr, step)
@@ -42,8 +41,6 @@ def main():
 
     f.close()
 
-    client.close()
-
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()
diff --git a/tools/mem_write.py b/tools/mem_write.py
index a868b397f..4bdaa4678 100755
--- a/tools/mem_write.py
+++ b/tools/mem_write.py
@@ -1,30 +1,27 @@
 #!/usr/bin/env python3
 from trezorlib.debuglink import DebugLink
-from trezorlib.client import TrezorClient
 from trezorlib.transport import enumerate_devices
 import sys
 
 
-def main():
-    # List all debuggable TREZORs
-    devices = [device for device in enumerate_devices() if hasattr(device, 'find_debug')]
-
-    # Check whether we found any
-    if len(devices) == 0:
-        print('No TREZOR found')
-        return
+def find_debug():
+    for device in enumerate_devices():
+        try:
+            debug_transport = device.find_debug()
+            debug = DebugLink(debug_transport, auto_interact=False)
+            debug.open()
+            return debug
+        except Exception:
+            continue
+    else:
+        print("No suitable Trezor device found")
+        sys.exit(1)
 
-    # Use first connected device
-    transport = devices[0]
-    debug_transport = devices[0].find_debug()
-
-    # Creates object for manipulating TREZOR
-    client = TrezorClient(transport)
-    debug = DebugLink(debug_transport)
 
+def main():
+    debug = find_debug()
     debug.memory_write(int(sys.argv[1], 16), bytes.fromhex(sys.argv[2]), flash=True)
-    client.close()
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()