Use toolbar to select channel

nrf802154_sniffer/nrf802154_sniffer.py

@@ -69,6 +69,19 @@ class Nrf802154Sniffer(object):
 
     # Helpers for Wireshark argument parsing.
     CTRL_ARG_CHANNEL = 0
+    CTRL_ARG_NONE = 255
+
+    # Helpers for Wireshark commands parsing.
+    CTRL_CMD_INITIALIZED = 0
+    CTRL_CMD_SET         = 1
+    CTRL_CMD_ADD         = 2
+    CTRL_CMD_REMOVE      = 3
+    CTRL_CMD_ENABLE      = 4
+    CTRL_CMD_DISABLE     = 5
+    CTRL_CMD_STATUSBAR   = 6
+    CTRL_CMD_INFORMATION = 7
+    CTRL_CMD_WARNING     = 8
+    CTRL_CMD_ERROR       = 9
 
     # Pattern for packets being printed over serial.
     RCV_REGEX = 'received:\s+([0-9a-fA-F]+)\s+power:\s+(-?\d+)\s+lqi:\s+(\d+)\s+time:\s+(-?\d+)'
@@ -78,7 +91,10 @@ class Nrf802154Sniffer(object):
     def __init__(self):
         self.serial = None
         self.serial_queue = Queue.Queue()
+        self.control_queue = Queue.Queue()
         self.running = threading.Event()
+        self.initialized = threading.Event()
+        self.initialized.clear()
         self.setup_done = threading.Event()
         self.setup_done.clear()
         self.logger = logging.getLogger(__name__)
@@ -135,6 +151,17 @@ def correct_time(self, sniffer_timestamp):
 
             return sniffer_timestamp + overflow_count * self.TIMER_MAX
 
+    def update_channel(self, channel):
+        """
+        Function for updating sniffing channel after sniffing started
+
+        :param channel: channel number (11-26)
+        """
+        self.channel = channel
+        while self.running.is_set() and not self.setup_done.is_set():
+            time.sleep(0.1)
+        self.serial_queue.put(b'channel ' + str(self.channel).encode())
+
     def stop_sig_handler(self, *args, **kwargs):
         """
         Function responsible for stopping the sniffer firmware and closing all threads.
@@ -280,14 +307,56 @@ def control_read(fn):
     def control_reader(self, fifo):
         """
         Thread responsible for reading wireshark commands (read from fifo).
-        Related to not-yet-implemented wireshark toolbar features.
         """
         with open(fifo, 'rb', 0) as fn:
             arg = 0
             while arg != None:
                 arg, typ, payload = Nrf802154Sniffer.control_read(fn)
+
+                if typ == Nrf802154Sniffer.CTRL_CMD_INITIALIZED:
+                    self.initialized.set()
+                elif arg == Nrf802154Sniffer.CTRL_ARG_CHANNEL and typ == Nrf802154Sniffer.CTRL_CMD_SET and payload:
+                    self.update_channel(int(payload))
+
             self.stop_sig_handler()
 
+    def control_send(self, arg, typ, payload):
+        """
+        Function responsible for preparing Wireshark command and putting it to send queue.
+        """
+        packet = bytearray()
+        packet += struct.pack('>sBHBB', b'T', 0, len(payload) + 2, arg, typ)
+        if sys.version_info[0] >= 3 and isinstance(payload, str):
+            packet += payload.encode('utf-8')
+        else:
+            packet += payload
+        self.control_queue.put(packet)
+
+    def control_write(self, fn):
+        """
+        Function responsible for writing single command from command queue to Wireshark configuration fifo.
+        """
+        command = self.control_queue.get(block=True, timeout=1)
+        try:
+            fn.write(command)
+        except IOError:
+            self.logger.error("Cannot write to {}".format(self))
+            self.running.clear()
+
+    def control_writer(self, fifo):
+        """
+        Thread responsible for writing wireshark commands (write to fifo).
+        """
+        while self.running.is_set() and not self.initialized.is_set():
+            time.sleep(.1)  # Wait for initial control values
+
+        with open(fifo, 'wb', 0) as fn:
+            while self.running.is_set():
+                try:
+                    self.control_write(fn)
+                except Queue.Empty:
+                    pass
+
     def serial_write(self):
         """
         Function responsible for sending commands to serial port.
@@ -317,7 +386,7 @@ def serial_writer(self):
             except Queue.Empty:
                 break
 
-    def serial_reader(self, dev, channel, queue):
+    def serial_reader(self, dev, queue):
         """
         Thread responsible for reading from serial port, parsing the output and storing parsed packets into queue.
         """
@@ -338,7 +407,7 @@ def serial_reader(self, dev, channel, queue):
             init_cmd = []
             init_cmd.append(b'')
             init_cmd.append(b'sleep')
-            init_cmd.append(b'channel ' + bytes(str(channel).encode()))
+            init_cmd.append(b'channel ' + bytes(str(self.channel).encode()))
             for cmd in init_cmd:
                 self.serial_queue.put(cmd)
 
@@ -349,6 +418,7 @@ def serial_reader(self, dev, channel, queue):
                 msg = "{} did not reply properly to setup commands. Please re-plug the device and make sure firmware is correct. " \
                         "Recieved: {}\n".format(self, init_res)
                 self.logger.error(msg)
+                self.control_send(Nrf802154Sniffer.CTRL_ARG_NONE, Nrf802154Sniffer.CTRL_CMD_ERROR, msg.encode())
 
             self.serial_queue.put(b'receive')
             self.setup_done.set()
@@ -359,7 +429,7 @@ def serial_reader(self, dev, channel, queue):
                 ch = self.serial.read()
                 if ch == b'':
                     continue
-                elif ch != '\n':
+                elif ch != b'\n' and ch != '\n':
                     buf += ch
                 else:
                     m = re.search(self.RCV_REGEX, str(buf))
@@ -368,7 +438,7 @@ def serial_reader(self, dev, channel, queue):
                         rssi = int(m.group(2))
                         lqi = int(m.group(3))
                         timestamp = int(m.group(4)) & 0xffffffff
-                        channel = int(channel)
+                        channel = int(self.channel)
                         queue.put(self.pcap_packet(packet, channel, rssi, lqi, self.correct_time(timestamp)))
                     buf = b''
 
@@ -412,17 +482,25 @@ def extcap_capture(self, fifo, dev, channel, control_in=None, control_out=None):
         self.dev = dev
         self.running.set()
 
-        # TODO: Add toolbar with channel selector (channel per interface?)
+        if control_out:
+            self.threads.append(threading.Thread(target=self.control_writer, args=(control_out,)))
+
         if control_in:
             self.threads.append(threading.Thread(target=self.control_reader, args=(control_in,)))
 
-        self.threads.append(threading.Thread(target=self.serial_reader, args=(self.dev, self.channel, packet_queue), name="serial_reader"))
+        self.threads.append(threading.Thread(target=self.serial_reader, args=(self.dev, packet_queue), name="serial_reader"))
         self.threads.append(threading.Thread(target=self.serial_writer, name="serial_writer"))
         self.threads.append(threading.Thread(target=self.fifo_writer, args=(fifo, packet_queue), name="fifo_writer"))
 
         for thread in self.threads:
             thread.start()
 
+        while self.running.is_set() and not self.initialized.is_set():
+            time.sleep(0.1)
+
+        time.sleep(0.1)
+        self.control_send(Nrf802154Sniffer.CTRL_ARG_CHANNEL, Nrf802154Sniffer.CTRL_CMD_SET, str(self.channel).encode())
+
         while is_standalone and self.running.is_set():
             time.sleep(1)