wificontrol: Add ReconnectWorker class

wificontrol/wifimonitor.py

@@ -4,6 +4,8 @@
 import signal
 from wificontrol import WiFiControl
 from reachstatus import StateClient
+from daemon_tree import DaemonTreeObj, DaemonTreeError
+from wifireconnect import WORKER_NAME
 
 try:
     from gi.repository import GObject
@@ -60,19 +62,29 @@ def __init__(self):
         self.current_state = None
         self.current_ssid = None
 
-        self._initialize()
+        self.reconnect_worker = DaemonTreeObj(WORKER_NAME)
+        self.disconnect_reason = None
 
-    def _initialize(self):
-        try:
-            systemd_obj = self.bus.get_object(SYSTEMD_DBUS_SERVICE,
-                                              SYSTEMD_DBUS_OPATH)
-            self.sysd_manager = dbus.Interface(systemd_obj,
-                                               dbus_interface=SYSTEMD_MANAGER_DBUS_IFACE)
-            self.sysd_manager.Subscribe()
+        self.register_callback(self.CLIENT_STATE, self._check_current_ssid)
+        self.register_callback(self.CLIENT_STATE, self._stop_reconnect_worker)
+
+        self.register_callback(self.HOST_STATE, self._clear_ssid)
+        self.register_callback(self.HOST_STATE, self._stop_reconnect_worker)
+
+        self.register_callback(self.OFF_STATE, self._check_disconnect_reason)
 
+        try:
+            self._initialize()
         except dbus.exceptions.DBusException as error:
             raise WiFiMonitorError(error)
 
+    def _initialize(self):
+        systemd_obj = self.bus.get_object(SYSTEMD_DBUS_SERVICE,
+                                          SYSTEMD_DBUS_OPATH)
+        self.sysd_manager = dbus.Interface(systemd_obj,
+                                           dbus_interface=SYSTEMD_MANAGER_DBUS_IFACE)
+        self.sysd_manager.Subscribe()
+
         self.bus.add_signal_receiver(self._wpa_props_changed,
                                      dbus_interface=WPAS_INTERFACE_DBUS_IFACE,
                                      signal_name="PropertiesChanged",
@@ -90,7 +102,7 @@ def _set_initial_state(self):
         self._process_new_state(state)
 
     def _host_props_changed(self, *args):
-        unit, props, etc = args
+        _, props, _ = args
         active_state = props.get('ActiveState')
         sub_state = props.get('SubState')
 
@@ -99,9 +111,10 @@ def _host_props_changed(self, *args):
 
     def _wpa_props_changed(self, props):
         state = props.get('State')
-        disconnect = props.get('DisconnectReason')
+        disconnect = props.get('DisconnectReason', None)
 
-        if disconnect:
+        if disconnect is not None:
+            self.disconnect_reason = disconnect
             state = 'disconnected'
 
         if state:
@@ -111,15 +124,8 @@ def _process_new_state(self, state):
         state = self.STATES.get(state)
         if state and self.current_state != state:
             self.current_state = state
-            self._execute_state_event(state)
             self._execute_callbacks(state)
 
-    def _execute_state_event(self, state):
-        if state == self.CLIENT_STATE:
-            self._check_current_ssid()
-        elif state == self.HOST_STATE:
-            self._clear_ssid()
-
     def _check_current_ssid(self):
         event = self.REVERT_EVENT
 
@@ -133,8 +139,8 @@ def _ssid_updated(self):
 
         try:
             ssid = status['ssid']
-        except (KeyError, TypeError):
-            return False
+        except (KeyError, TypeError) as error:
+            raise WiFiMonitorError(error)
 
         if self.current_ssid != ssid:
             self.current_ssid = ssid
@@ -145,6 +151,23 @@ def _ssid_updated(self):
     def _clear_ssid(self):
         self.current_ssid = None
 
+    def _check_disconnect_reason(self):
+        if self.disconnect_reason == 0:
+            self._start_reconnect_worker()
+        self.disconnect_reason = None
+
+    def _start_reconnect_worker(self):
+        try:
+            self.reconnect_worker.call('start_reconnection', args=(self.current_ssid,))
+        except DaemonTreeError as error:
+            raise WiFiMonitorError(error)
+
+    def _stop_reconnect_worker(self):
+        try:
+            self.reconnect_worker.call('stop_reconnection')
+        except DaemonTreeError as error:
+            raise WiFiMonitorError(error)
+
     def register_callback(self, msg, callback, args=()):
         if msg not in self.callbacks:
             self.callbacks[msg] = []
@@ -156,7 +179,10 @@ def _execute_callbacks(self, msg):
         if callbacks:
             for callback in callbacks:
                 callback, args = callback
-                callback(*args)
+                try:
+                    callback(*args)
+                except Exception as error:
+                    raise WiFiMonitorError(error)
 
     def run(self):
         self._mainloop.run()

wificontrol/wifireconnect.py

@@ -0,0 +1,75 @@
+import dbus
+import signal
+import threading
+import logging
+from daemon_tree import DaemonTreeSvr
+from wificontrol import WiFiControl
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+
+WORKER_NAME = 'wifi_reconnect'
+
+
+class ReconnectWorker(object):
+    TIMEOUT = 10
+
+    def __init__(self):
+        self.manager = WiFiControl()
+        self.interrupt = threading.Event()
+        self.worker = None
+
+    def start_reconnection(self, ssid):
+        if self.worker is None:
+            self.worker = threading.Thread(target=self._reconnect, args=(ssid,))
+            self.worker.start()
+
+    def _reconnect(self, ssid):
+        self.interrupt.clear()
+        self.interrupt.wait(self.TIMEOUT)
+
+        while not self.interrupt.is_set():
+            try:
+                self.manager.scan()
+                scan_results = self.manager.get_scan_results()
+
+                scanned_ssids = [net['ssid'] for net in scan_results]
+
+                if ssid in scanned_ssids:
+                    network = {'ssid': ssid}
+                    self.manager.start_connecting(network, callback=self._callback)
+            except dbus.exceptions.DBusException as error:
+                logger.error(error)
+
+            self.interrupt.wait(self.TIMEOUT)
+
+    def _callback(self, result=None):
+        if result:
+            self.stop_reconnection()
+
+    def stop_reconnection(self):
+        self.interrupt.set()
+        if self.worker:
+            self.worker.join()
+            self.worker = None
+
+
+def main():
+    def handler(signum, frame):
+        reconnect_svr.cancel()
+
+    reconnect_worker = ReconnectWorker()
+    reconnect_svr = DaemonTreeSvr(name=WORKER_NAME)
+
+    reconnect_svr.register(reconnect_worker.start_reconnection)
+    reconnect_svr.register(reconnect_worker.stop_reconnection)
+
+    signal.signal(signal.SIGINT, handler)
+    signal.signal(signal.SIGTERM, handler)
+
+    reconnect_svr.run()
+    reconnect_svr.shutdown()
+
+
+if __name__ == '__main__':
+    main()

wificontrol/wpasupplicant.py

@@ -32,7 +32,8 @@
 
 
 class WpaSupplicant(WiFi):
-    wpas_control = lambda self, action: "systemctl {} wpa_supplicant.service && sleep 2".format(action)
+    wpas_control = lambda self, action: "systemctl {} wpa_supplicant.service && sleep 2".format(
+        action)
 
     def __init__(self, interface,
                  wpas_config="/etc/wpa_supplicant/wpa_supplicant.conf",
@@ -57,12 +58,15 @@ def __init__(self, interface,
         self.connection_timer = None
         self.break_event = Event()
 
-        self.started = lambda: self.sysdmanager.is_active(
+    def started(self):
+        wpa_supplicant_started = self.sysdmanager.is_active(
             "wpa_supplicant.service")
 
-        if self.started():
+        if wpa_supplicant_started:
             self.wpa_supplicant_interface.initialize()
 
+        return wpa_supplicant_started
+
     def start(self):
         self.execute_command(self.wpas_control("start"))
         self.wpa_supplicant_interface.initialize()