From 5bc6b530d89045f25e246d471fbb7deb70e7f10c Mon Sep 17 00:00:00 2001
From: Robin Krens <robin@robinkrens.nl>
Date: Wed, 14 Feb 2024 15:27:27 +0100
Subject: [PATCH] RAFlasher.py: move inquire connection to RAConnect class

---
 src/RAConnect.py | 18 +++++++++++++++++-
 src/RAFlasher.py | 32 ++++++--------------------------
 2 files changed, 23 insertions(+), 27 deletions(-)

diff --git a/src/RAConnect.py b/src/RAConnect.py
index f753d24..8405dae 100644
--- a/src/RAConnect.py
+++ b/src/RAConnect.py
@@ -2,6 +2,7 @@ import sys
 import time
 import usb.core
 import usb.util
+from RAPacker import *
 
 MAX_TRANSFER_SIZE = 2048 + 6 # include header and footer
 
@@ -16,7 +17,11 @@ class RAConnect:
         self.dev = None
         self.rx_ep = None
         self.tx_ep = None
+
         self.find_device()
+        status_conn = self.inquire_connection()
+        if not status_conn:
+            self.confirm_connection()
 
     def find_device(self):
         self.dev = usb.core.find(idVendor=self.vendor_id, idProduct=self.product_id)
@@ -38,7 +43,18 @@ class RAConnect:
                         self.tx_ep = ep
             return True
 
-        raise ValueError("Device does not have a serial interface")
+        raise ValueError("Device does not have a CDC interface")
+
+
+    def inquire_connection(self):
+        packed = pack_pkt(INQ_CMD, "")
+        self.send_data(packed)
+        info = self.recv_data(7)
+        if info == bytearray(b'\x00') or info == bytearray(b''):
+            return False
+        msg = unpack_pkt(info)
+        #print("Connection already established")
+        return True
 
     def establish_connection(self):
         for i in range(self.max_tries):
diff --git a/src/RAFlasher.py b/src/RAFlasher.py
index e71aee7..025bee4 100644
--- a/src/RAFlasher.py
+++ b/src/RAFlasher.py
@@ -8,6 +8,8 @@ from RAConnect import *
 from RAPacker import *
 
 SECTOR_SIZE = 2048
+VENDOR_ID = 0x045B
+PRODUCT_ID = 0x0261
 
 def int_to_hex_list(num):
     hex_string = hex(num)[2:].upper()  # convert to hex string
@@ -40,16 +42,6 @@ def set_size_boundaries(start_addr, size):
 
     return (start_addr, end_addr)
 
-def inquire_connection(dev):
-    packed = pack_pkt(INQ_CMD, "")
-    dev.send_data(packed)
-    info = dev.recv_data(7)
-    if info == bytearray(b'\x00') or info == bytearray(b''):
-        return False
-    msg = unpack_pkt(info)
-    #print("Connection already established")
-    return True
-
 def get_area_info(dev):
     for i in [0,1,2]:
         packed = pack_pkt(ARE_CMD, [str(i)])
@@ -206,28 +198,16 @@ def main():
     args = parser.parse_args()
 
     if args.command == "write":
-        dev = RAConnect(vendor_id=0x045B, product_id=0x0261)
-        status_con = inquire_connection(dev)
-        if not status_con:
-            dev.confirm_connection()
+        dev = RAConnect(VENDOR_ID, PRODUCT_ID)
         write_img(dev, args.file_name, args.start_address, args.size, args.verify)
     elif args.command == "read":
-        dev = RAConnect(vendor_id=0x045B, product_id=0x0261)
-        status_con = inquire_connection(dev)
-        if not status_con:
-            dev.confirm_connection()
+        dev = RAConnect(VENDOR_ID, PRODUCT_ID)
         read_img(dev, args.file_name, args.start_address, args.size)
     elif args.command == "erase":
-        dev = RAConnect(vendor_id=0x045B, product_id=0x0261)
-        status_con = inquire_connection(dev)
-        if not status_con:
-            dev.confirm_connection()
+        dev = RAConnect(VENDOR_ID, PRODUCT_ID)
         erase_chip(dev, args.start_address, args.size)
     elif args.command == "info":
-        dev = RAConnect(vendor_id=0x045B, product_id=0x0261)
-        status_con = inquire_connection(dev)
-        if not status_con:
-            dev.confirm_connection()
+        dev = RAConnect(VENDOR_ID, PRODUCT_ID)
         get_dev_info(dev)
         get_area_info(dev)
     else:
-- 
2.7.4