From 95830698789f1641c5e8ab68b96e0309387c077b Mon Sep 17 00:00:00 2001
From: Dmytro Yaroshenko <73843436+o-murphy@users.noreply.github.com>
Date: Mon, 29 Apr 2024 17:22:16 +0300
Subject: [PATCH] trying to implement main.py

---
 pydfuutil/__main__.py | 31 +++++++++++++++++++------------
 pydfuutil/usb_dfu.py  | 27 ++++++++++++++-------------
 2 files changed, 33 insertions(+), 25 deletions(-)

diff --git a/pydfuutil/__main__.py b/pydfuutil/__main__.py
index cf1749a..1af78d6 100644
--- a/pydfuutil/__main__.py
+++ b/pydfuutil/__main__.py
@@ -543,7 +543,6 @@ def le16_to_cpu(data):
     return int.from_bytes(data, byteorder='little')
 
 
-
 class IntOrBytes:
     def __init__(self, value):
         if isinstance(value, int):
@@ -621,7 +620,6 @@ def main() -> None:
     mode = Mode.NONE
     device_id_filter = None
 
-
     func_dfu_rt = USB_DFU_FUNC_DESCRIPTOR.parse(bytes(USB_DFU_FUNC_DESCRIPTOR.sizeof()))
 
     if args.verbose:
@@ -770,6 +768,7 @@ def get_first_dfu_if(dif_: dfu.DfuIf, v: Any = None):
     # Transition from run-Time mode to DFU mode
 
     if not (_rt_dif.flags & dfu.Mode.IFF_DFU):
+
         # In the 'first round' during runtime mode, there can only be one
         # DFU Interface descriptor according to the DFU Spec.
 
@@ -797,22 +796,30 @@ def get_first_dfu_if(dif_: dfu.DfuIf, v: Any = None):
         if not quirks & QUIRK_POLLTIMEOUT:
             milli_sleep(status.bwPollTimeout)
 
-        if status.bState in [dfu.State.APP_IDLE, dfu.State.APP_DETACH]:
+        if status.bState in (dfu.State.APP_IDLE, dfu.State.APP_DETACH):
             print("Device really in Runtime Mode, send DFU "
                   "detach request...")
 
-            ret = dfu.detach(_rt_dif.dev, _rt_dif.interface, 1000)
-            print(ret)
-            if IntOrBytes(ret) < 0:
+            if IntOrBytes(dfu.detach(_rt_dif.dev, _rt_dif.interface, 1000)) < 0:
                 print("error detaching")
                 exit(1)
 
-            # print(bmAttributes.USB_DFU_WILL_DETACH.__dir__())
-            # print(func_dfu_rt.bmAttributes)
-            # # if func_dfu_rt.bmAttributes & bmAttributes.USB_DFU_WILL_DETACH:
-            # #     print('continue')
-
-
+            if func_dfu_rt.bmAttributes & bmAttributes.USB_DFU_WILL_DETACH:
+                print("Device will detach and reattach...")
+            else:
+                print("Resetting USB...\n")
+                try:
+                    _rt_dif.dev.reset()
+                except usb.core.USBError as exc:
+                    print("error resetting after detach")
+            milli_sleep(2000)
+        elif status.bState == dfu.State.DFU_ERROR:
+            print("dfuERROR, clearing status")
+            if IntOrBytes(dfu.clear_status(_rt_dif.dev, _rt_dif.interface)) < 0:
+                print("error detaching")
+                exit(1)
+        else:
+            print("WARNING: Runtime device already in DFU state ?!?")
 
 
 
diff --git a/pydfuutil/usb_dfu.py b/pydfuutil/usb_dfu.py
index fcbdc15..c6be43d 100644
--- a/pydfuutil/usb_dfu.py
+++ b/pydfuutil/usb_dfu.py
@@ -14,24 +14,25 @@
 
 USB_DT_DFU = 0x21
 
-# class bmAttributes(IntEnum):
-#     USB_DFU_CAN_DOWNLOAD = 0x1
-#     USB_DFU_CAN_UPLOAD = 0x2
-#     USB_DFU_MANIFEST_TOL = 0x3
-#     USB_DFU_WILL_DETACH = 0x4
+class bmAttributes(IntEnum):
+    USB_DFU_CAN_DOWNLOAD = 0x1
+    USB_DFU_CAN_UPLOAD = 0x2
+    USB_DFU_MANIFEST_TOL = 0x3
+    USB_DFU_WILL_DETACH = 0x4
 
-bmAttributes = FlagsEnum(
-    Byte,
-    USB_DFU_CAN_DOWNLOAD=0x1,  # is support updates
-    USB_DFU_CAN_UPLOAD=0x2,  # is prog warranty ok
-    USB_DFU_MANIFEST_TOL=0x4,
-    USB_DFU_WILL_DETACH=0x8,
-)
+# bmAttributes = FlagsEnum(
+#     Byte,
+#     USB_DFU_CAN_DOWNLOAD=0x1,  # is support updates
+#     USB_DFU_CAN_UPLOAD=0x2,  # is prog warranty ok
+#     USB_DFU_MANIFEST_TOL=0x4,
+#     USB_DFU_WILL_DETACH=0x8,
+# )
 
 USB_DFU_FUNC_DESCRIPTOR = Struct(
     bLength=Int8ul,
     bDescriptorType=Int8ul,
-    bmAttributes=bmAttributes,
+    # bmAttributes=bmAttributes,
+    bmAttributes=Int8ul,
     wDetachTimeOut=Int16ul,
     wTransferSize=Int16ul,
     bcdDFUVersion=Int16ul,