2nd batch update

drivers/NRF52_dongle.py

@@ -3,6 +3,7 @@
 import sys
 
 import serial
+import serial.tools.list_ports
 
 sys.path.insert(0, os.getcwd() + '../libs')
 from colorama import Fore
@@ -11,9 +12,14 @@
 
 # USB Serial commands
 NRF52_CMD_DATA = b'\xA7'
+NRF52_CMD_DATA_TX = b'\xBB'
 NRF52_CMD_CHECKSUM_ERROR = b'\xA8'
 NRF52_CMD_CONFIG_AUTO_EMPTY_PDU = b'\xA9'
 NRF52_CMD_CONFIG_ACK = b'\xAA'
+NRF52_CMD_CONFIG_LOG_TX = b'\xCC'
+NRF52_CMD_CONFIG_NESNSN = b'\xAD'
+NRF52_CMD_CONFIG_NESN = b'\xAE'
+NRF52_CMD_CONFIG_SN = b'\xAF'
 NRF52_CMD_BOOTLOADER_SEQ1 = b'\xA6'
 NRF52_CMD_BOOTLOADER_SEQ2 = b'\xC7'
 NRF52_CMD_LOG = b'\x7F'
@@ -27,9 +33,22 @@ class NRF52Dongle:
     event_counter = 0
     packets_buffer = []
     pcap_filename = None
+    pcap_tx_handover = False
+    sent_pkt = None
 
     # Constructor ------------------------------------
-    def __init__(self, port_name, baudrate, debug=False, logs=True, logs_pcap=False, pcap_filename=None):
+    def __init__(self, port_name=None, baudrate=115200, debug=False, logs=True, logs_pcap=False, pcap_filename=None):
+
+        if port_name is None:
+            found = False
+            ports = serial.tools.list_ports.comports()
+            for port in ports:
+                if 'Bluefruit nRF52840' in port.description:
+                    port_name = port.device
+                    found = True
+            if not found:
+                print(Fore.RED + 'nRF52840 was not found')
+
         self.serial = serial.Serial(port_name, baudrate, timeout=1)
         self.logs_pcap = logs_pcap
         self.n_log = logs
@@ -39,6 +58,8 @@ def __init__(self, port_name, baudrate, debug=False, logs=True, logs_pcap=False,
         else:
             self.pcap_filename = pcap_filename
 
+        self.set_log_tx(0)
+
         if self.n_debug:
             print('NRF52 Dongle: Instance started')
 
@@ -64,17 +85,17 @@ def raw_send(self, pkt):
 
         return data
 
-    def send(self, scapy_pkt, print_tx=True):
+    def send(self, scapy_pkt, print_tx=True, force_pcap_save=False):
         self.raw_send(raw(scapy_pkt))
-        if self.logs_pcap:
+        if self.logs_pcap and (self.pcap_tx_handover is 0 or force_pcap_save):
             self.packets_buffer.append(NORDIC_BLE(board=75, protocol=2, flags=0x3) / scapy_pkt)
         if print_tx:
             print(Fore.CYAN + "TX ---> " + scapy_pkt.summary()[7:])
 
     def raw_receive(self):
         c = self.serial.read(1)
         # Receive BLE adv or channel packets
-        if c == NRF52_CMD_DATA:
+        if c == NRF52_CMD_DATA or c == NRF52_CMD_DATA_TX:
             lb = ord(self.serial.read(1))
             hb = ord(self.serial.read(1))
             sz = lb | (hb << 8)
@@ -87,12 +108,21 @@ def raw_receive(self):
                 # If the data received is correct
                 self.event_counter = evt_counter
 
+                if c == NRF52_CMD_DATA_TX:
+                    self.sent_pkt = data
+                    n_flags = 0x03
+                    ret_data = None
+                else:  # Received packets
+                    n_flags = 0x01
+                    ret_data = data
+
+                if self.logs_pcap is True and data != None:
+                    self.packets_buffer.append(NORDIC_BLE(board=75, protocol=2, flags=n_flags) / BTLE(data))
+
                 if self.n_debug:
                     print("Hex: " + binascii.hexlify(data).upper())
-                if self.logs_pcap is True and data != None:
-                    self.packets_buffer.append(NORDIC_BLE(board=75, protocol=2, flags=0x01) / BTLE(data))
 
-                return data
+                return ret_data
         # Receive logs from dongle
         elif c == NRF52_CMD_LOG:
             lb = ord(self.serial.read(1))
@@ -103,5 +133,32 @@ def raw_receive(self):
                 print(data)
         elif c == NRF52_CMD_CHECKSUM_ERROR:
             print(Fore.RED + "NRF52_CMD_CHECKSUM_ERROR")
-            sys.exit(0)
+
         return None
+
+    # Set the initial value of NESN and SN
+    def set_nesnsn(self, value):
+        # 0b01 -> set NESN, 0b10 -> set SN
+        data = NRF52_CMD_CONFIG_NESNSN + bytearray([value])
+        self.serial.write(data)
+
+    # Set the initial value of NESN
+    def set_nesn(self, value):
+        data = NRF52_CMD_CONFIG_NESN + bytearray([value])
+        self.serial.write(data)
+
+    # Set the initial value of SN
+    def set_sn(self, value):
+        data = NRF52_CMD_CONFIG_SN + bytearray([value])
+        self.serial.write(data)
+
+    def get_tx_packet(self):
+        pkt = self.sent_pkt
+        self.sent_pkt = None
+        return pkt
+
+    # Makes the dongle send back the processed tx packet
+    def set_log_tx(self, value):
+        data = NRF52_CMD_CONFIG_LOG_TX + bytearray([value])
+        self.serial.write(data)
+        self.pcap_tx_handover = value