6510b2ec6b
Use f-strings Avoid unnecessary casts to str()bool()int() Ensure file closure Merge nested ifs Simplify for-assigns-appends with comprehensions and internal functions Avoid shadowing internal function names Initialize variables Return value directly instead of assign then return Make some methods static
474 lines
15 KiB
Python
474 lines
15 KiB
Python
import os
|
|
import logging
|
|
import re
|
|
import socket
|
|
import threading
|
|
|
|
from threading import Thread
|
|
|
|
from queue import Queue
|
|
|
|
import gi
|
|
|
|
gi.require_version("Gtk", "3.0")
|
|
from gi.repository import GLib, Gdk
|
|
|
|
|
|
class WifiManager:
|
|
networks_in_supplicant = []
|
|
connected = False
|
|
_stop_loop = False
|
|
thread = None
|
|
|
|
def __init__(self, interface, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.loop = None
|
|
self._poll_task = None
|
|
self._scanning = False
|
|
self._callbacks = {
|
|
"connected": [],
|
|
"connecting_status": [],
|
|
"scan_results": []
|
|
}
|
|
self._stop_loop = False
|
|
self.connected = False
|
|
self.connected_ssid = None
|
|
self.connecting_info = []
|
|
self.event = threading.Event()
|
|
self.initialized = False
|
|
self.interface = interface
|
|
self.networks = {}
|
|
self.supplicant_networks = {}
|
|
self.queue = Queue()
|
|
self.tasks = []
|
|
self.timeout = None
|
|
self.scan_time = 0
|
|
|
|
ks_socket_file = "/tmp/.KS_wpa_supplicant"
|
|
if os.path.exists(ks_socket_file):
|
|
os.remove(ks_socket_file)
|
|
|
|
try:
|
|
self.soc = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
|
|
self.soc.bind(ks_socket_file)
|
|
self.soc.connect(f"/var/run/wpa_supplicant/{interface}")
|
|
except Exception as e:
|
|
logging.critical(e, exc_info=True)
|
|
logging.error(f"Error connecting to wifi socket: {interface}")
|
|
return
|
|
|
|
self.wpa_thread = WpaSocket(self, self.queue, self.callback)
|
|
self.wpa_thread.start()
|
|
self.initialized = True
|
|
|
|
self.wpa_cli("ATTACH", False)
|
|
self.wpa_cli("SCAN", False)
|
|
GLib.idle_add(self.read_wpa_supplicant)
|
|
self.timeout = GLib.timeout_add_seconds(180, self.rescan)
|
|
|
|
def add_callback(self, name, callback):
|
|
if name in self._callbacks and callback not in self._callbacks[name]:
|
|
self._callbacks[name].append(callback)
|
|
|
|
def add_network(self, ssid, psk):
|
|
for netid in list(self.supplicant_networks):
|
|
if self.supplicant_networks[netid]['ssid'] == ssid:
|
|
# Modify network
|
|
return
|
|
|
|
# TODO: Add wpa_cli error checking
|
|
network_id = self.wpa_cli("ADD_NETWORK")
|
|
commands = [
|
|
f'ENABLE_NETWORK {network_id}',
|
|
'SET_NETWORK %s ssid "%s"' % (network_id, ssid.replace('"', '\"')),
|
|
'SET_NETWORK %s psk "%s"' % (network_id, psk.replace('"', '\"'))
|
|
]
|
|
|
|
self.wpa_cli_batch(commands)
|
|
|
|
self.read_wpa_supplicant()
|
|
netid = None
|
|
for i in list(self.supplicant_networks):
|
|
if self.supplicant_networks[i]['ssid'] == ssid:
|
|
netid = i
|
|
break
|
|
|
|
if netid is None:
|
|
logging.info("Error adding network")
|
|
return False
|
|
|
|
self.save_wpa_conf()
|
|
return True
|
|
|
|
def callback(self, cb_type, msg):
|
|
if cb_type in self._callbacks:
|
|
for cb in self._callbacks[cb_type]:
|
|
Gdk.threads_add_idle(
|
|
GLib.PRIORITY_DEFAULT_IDLE,
|
|
cb,
|
|
msg)
|
|
|
|
def connect(self, ssid):
|
|
netid = None
|
|
for nid, net in self.supplicant_networks.items():
|
|
if net['ssid'] == ssid:
|
|
netid = nid
|
|
break
|
|
|
|
if netid is None:
|
|
logging.info("Wifi network is not defined in wpa_supplicant")
|
|
return False
|
|
|
|
logging.info(f"Attempting to connect to wifi: {netid}")
|
|
self.connecting_info = [f"Attempting to connect to {ssid}"]
|
|
self.wpa_cli(f"SELECT_NETWORK {id}")
|
|
self.save_wpa_conf()
|
|
|
|
def delete_network(self, ssid):
|
|
netid = None
|
|
for i in list(self.supplicant_networks):
|
|
if self.supplicant_networks[i]['ssid'] == ssid:
|
|
netid = i
|
|
break
|
|
|
|
if netid is None:
|
|
logging.debug("Unable to find network in wpa_supplicant")
|
|
return
|
|
self.wpa_cli(f"REMOVE_NETWORK {netid}")
|
|
|
|
for netid in list(self.supplicant_networks):
|
|
if self.supplicant_networks[netid]['ssid'] == ssid:
|
|
del self.supplicant_networks[netid]
|
|
break
|
|
|
|
self.save_wpa_conf()
|
|
|
|
def get_connected_ssid(self):
|
|
return self.connected_ssid
|
|
|
|
def get_current_wifi(self):
|
|
con_ssid = os.popen("sudo iwgetid -r").read().strip()
|
|
con_bssid = os.popen("sudo iwgetid -r -a").read().strip()
|
|
# wpa_cli status output is unstable use it as backup only
|
|
status = self.wpa_cli("STATUS").split('\n')
|
|
variables = {}
|
|
for line in status:
|
|
arr = line.split('=')
|
|
variables[arr[0]] = "=".join(arr[1:])
|
|
prev_ssid = self.connected_ssid
|
|
|
|
if con_ssid != "":
|
|
self.connected = True
|
|
self.connected_ssid = con_ssid
|
|
for ssid, val in self.networks.items():
|
|
self.networks[ssid]['connected'] = ssid == con_ssid
|
|
if prev_ssid != self.connected_ssid:
|
|
for cb in self._callbacks['connected']:
|
|
Gdk.threads_add_idle(
|
|
GLib.PRIORITY_DEFAULT_IDLE,
|
|
cb, self.connected_ssid, prev_ssid)
|
|
return [con_ssid, con_bssid]
|
|
elif "ssid" in variables and "bssid" in variables:
|
|
self.connected = True
|
|
self.connected_ssid = variables['ssid']
|
|
for ssid, val in self.networks.items():
|
|
self.networks[ssid]['connected'] = ssid == variables['ssid']
|
|
if prev_ssid != self.connected_ssid:
|
|
for cb in self._callbacks['connected']:
|
|
Gdk.threads_add_idle(
|
|
GLib.PRIORITY_DEFAULT_IDLE,
|
|
cb, self.connected_ssid, prev_ssid)
|
|
return [variables['ssid'], variables['bssid']]
|
|
else:
|
|
logging.info("Resetting connected_ssid")
|
|
self.connected = False
|
|
self.connected_ssid = None
|
|
for ssid, val in self.networks.items():
|
|
self.networks[ssid]['connected'] = False
|
|
if prev_ssid != self.connected_ssid:
|
|
for cb in self._callbacks['connected']:
|
|
Gdk.threads_add_idle(
|
|
GLib.PRIORITY_DEFAULT_IDLE,
|
|
cb, self.connected_ssid, prev_ssid)
|
|
return None
|
|
|
|
def get_current_wifi_idle_add(self):
|
|
self.get_current_wifi()
|
|
return False
|
|
|
|
def get_network_info(self, ssid=None, mac=None):
|
|
if ssid is not None and ssid in self.networks:
|
|
return self.networks[ssid]
|
|
if mac is not None and ssid is None:
|
|
for net in self.networks:
|
|
if mac == net['mac']:
|
|
return net
|
|
return None
|
|
|
|
def get_networks(self):
|
|
return list(self.networks)
|
|
|
|
def get_supplicant_networks(self):
|
|
return self.supplicant_networks
|
|
|
|
def is_connected(self):
|
|
return self.connected
|
|
|
|
def is_initialized(self):
|
|
return self.initialized
|
|
|
|
def read_wpa_supplicant(self):
|
|
results = self.wpa_cli("LIST_NETWORKS").split('\n')
|
|
results.pop(0)
|
|
self.supplicant_networks = {}
|
|
self.networks_in_supplicant = []
|
|
for net in [n.split('\t') for n in results]:
|
|
self.supplicant_networks[net[0]] = {
|
|
"ssid": net[1],
|
|
"bssid": net[2],
|
|
"flags": net[3] if len(net) == 4 else ""
|
|
}
|
|
self.networks_in_supplicant.append(self.supplicant_networks[net[0]])
|
|
|
|
def remove_callback(self, name, callback):
|
|
if name in self._callbacks and callback in self._callbacks[name]:
|
|
self._callbacks[name].remove(callback)
|
|
|
|
def rescan(self):
|
|
self.wpa_cli("SCAN", False)
|
|
return True
|
|
|
|
def save_wpa_conf(self):
|
|
logging.info("Saving WPA config")
|
|
self.wpa_cli("SAVE_CONFIG")
|
|
|
|
def scan_results(self):
|
|
new_networks = []
|
|
deleted_networks = list(self.networks)
|
|
|
|
results = self.wpa_cli("SCAN_RESULTS").split('\n')
|
|
results.pop(0)
|
|
|
|
aps = []
|
|
for res in results:
|
|
match = re.match("^([a-f0-9:]+)\\s+([0-9]+)\\s+([\\-0-9]+)\\s+(\\S+)\\s+(.+)?", res)
|
|
if match:
|
|
net = {
|
|
"mac": match[1],
|
|
"channel": WifiChannels.lookup(match[2])[1],
|
|
"connected": False,
|
|
"configured": False,
|
|
"frequency": match[2],
|
|
"flags": match[4],
|
|
"signal_level_dBm": match[3],
|
|
"ssid": match[5]
|
|
}
|
|
|
|
if "WPA2" in net['flags']:
|
|
net['encryption'] = "WPA2"
|
|
elif "WPA" in net['flags']:
|
|
net['encryption'] = "WPA"
|
|
elif "WEP" in net['flags']:
|
|
net['encryption'] = "WEP"
|
|
else:
|
|
net['encryption'] = "off"
|
|
|
|
aps.append(net)
|
|
|
|
cur_info = self.get_current_wifi()
|
|
self.networks = {}
|
|
for ap in aps:
|
|
self.networks[ap['ssid']] = ap
|
|
if cur_info is not None and cur_info[0] == ap['ssid'] and cur_info[1].lower() == ap['mac'].lower():
|
|
self.networks[ap['ssid']]['connected'] = True
|
|
|
|
for net in list(self.networks):
|
|
if net in deleted_networks:
|
|
deleted_networks.remove(net)
|
|
else:
|
|
new_networks.append(net)
|
|
if new_networks or deleted_networks:
|
|
for cb in self._callbacks['scan_results']:
|
|
Gdk.threads_add_idle(
|
|
GLib.PRIORITY_DEFAULT_IDLE,
|
|
cb, new_networks, deleted_networks)
|
|
|
|
def wpa_cli(self, command, wait=True):
|
|
if wait is False:
|
|
self.wpa_thread.skip_command()
|
|
self.soc.send(command.encode())
|
|
if wait is True:
|
|
return self.queue.get()
|
|
|
|
def wpa_cli_batch(self, commands):
|
|
for cmd in commands:
|
|
self.wpa_cli(cmd)
|
|
|
|
|
|
class WpaSocket(Thread):
|
|
def __init__(self, wm, queue, callback):
|
|
super().__init__()
|
|
self.queue = queue
|
|
self.callback = callback
|
|
self.soc = wm.soc
|
|
self._stop_loop = False
|
|
self.skip_commands = 0
|
|
self.wm = wm
|
|
|
|
def run(self):
|
|
logging.debug("Setting up wifi event loop")
|
|
while self._stop_loop is False:
|
|
try:
|
|
msg = self.soc.recv(4096).decode().strip()
|
|
except Exception as e:
|
|
logging.critical(e, exc_info=True)
|
|
# TODO: Socket error
|
|
continue
|
|
if msg.startswith("<"):
|
|
if "CTRL-EVENT-SCAN-RESULTS" in msg:
|
|
Gdk.threads_add_idle(GLib.PRIORITY_DEFAULT_IDLE, self.wm.scan_results)
|
|
elif "CTRL-EVENT-DISCONNECTED" in msg:
|
|
self.callback("connecting_status", msg)
|
|
match = re.match('<3>CTRL-EVENT-DISCONNECTED bssid=(\\S+) reason=3 locally_generated=1', msg)
|
|
if match:
|
|
for net in self.wm.networks:
|
|
if self.wm.networks[net]['mac'] == match[1]:
|
|
self.wm.networks[net]['connected'] = False
|
|
break
|
|
elif "Trying to associate" in msg or "CTRL-EVENT-REGDOM-CHANGE" in msg:
|
|
self.callback("connecting_status", msg)
|
|
elif "CTRL-EVENT-CONNECTED" in msg:
|
|
Gdk.threads_add_idle(GLib.PRIORITY_DEFAULT_IDLE, self.wm.get_current_wifi_idle_add)
|
|
self.callback("connecting_status", msg)
|
|
elif self.skip_commands > 0:
|
|
self.skip_commands = self.skip_commands - 1
|
|
else:
|
|
self.queue.put(msg)
|
|
logging.info("Wifi event loop ended")
|
|
|
|
def skip_command(self):
|
|
self.skip_commands = self.skip_commands + 1
|
|
|
|
def stop(self):
|
|
self._stop_loop = True
|
|
|
|
|
|
class WifiChannels:
|
|
@staticmethod
|
|
def lookup(freq):
|
|
if freq == "2412":
|
|
return "2.4", "1"
|
|
if freq == "2417":
|
|
return "2.4", "2"
|
|
if freq == "2422":
|
|
return "2.4", "3"
|
|
if freq == "2427":
|
|
return "2.4", "4"
|
|
if freq == "2432":
|
|
return "2.4", "5"
|
|
if freq == "2437":
|
|
return "2.4", "6"
|
|
if freq == "2442":
|
|
return "2.4", "7"
|
|
if freq == "2447":
|
|
return "2.4", "8"
|
|
if freq == "2452":
|
|
return "2.4", "9"
|
|
if freq == "2457":
|
|
return "2.4", "10"
|
|
if freq == "2462":
|
|
return "2.4", "11"
|
|
if freq == "2467":
|
|
return "2.4", "12"
|
|
if freq == "2472":
|
|
return "2.4", "13"
|
|
if freq == "2484":
|
|
return "2.4", "14"
|
|
if freq == "5035":
|
|
return "5", "7"
|
|
if freq == "5040":
|
|
return "5", "8"
|
|
if freq == "5045":
|
|
return "5", "9"
|
|
if freq == "5055":
|
|
return "5", "11"
|
|
if freq == "5060":
|
|
return "5", "12"
|
|
if freq == "5080":
|
|
return "5", "16"
|
|
if freq == "5170":
|
|
return "5", "34"
|
|
if freq == "5180":
|
|
return "5", "36"
|
|
if freq == "5190":
|
|
return "5", "38"
|
|
if freq == "5200":
|
|
return "5", "40"
|
|
if freq == "5210":
|
|
return "5", "42"
|
|
if freq == "5220":
|
|
return "5", "44"
|
|
if freq == "5230":
|
|
return "5", "46"
|
|
if freq == "5240":
|
|
return "5", "48"
|
|
if freq == "5260":
|
|
return "5", "52"
|
|
if freq == "5280":
|
|
return "5", "56"
|
|
if freq == "5300":
|
|
return "5", "60"
|
|
if freq == "5320":
|
|
return "5", "64"
|
|
if freq == "5500":
|
|
return "5", "100"
|
|
if freq == "5520":
|
|
return "5", "104"
|
|
if freq == "5540":
|
|
return "5", "108"
|
|
if freq == "5560":
|
|
return "5", "112"
|
|
if freq == "5580":
|
|
return "5", "116"
|
|
if freq == "5600":
|
|
return "5", "120"
|
|
if freq == "5620":
|
|
return "5", "124"
|
|
if freq == "5640":
|
|
return "5", "128"
|
|
if freq == "5660":
|
|
return "5", "132"
|
|
if freq == "5680":
|
|
return "5", "136"
|
|
if freq == "5700":
|
|
return "5", "140"
|
|
if freq == "5720":
|
|
return "5", "144"
|
|
if freq == "5745":
|
|
return "5", "149"
|
|
if freq == "5765":
|
|
return "5", "153"
|
|
if freq == "5785":
|
|
return "5", "157"
|
|
if freq == "5805":
|
|
return "5", "161"
|
|
if freq == "5825":
|
|
return "5", "165"
|
|
if freq == "4915":
|
|
return "5", "183"
|
|
if freq == "4920":
|
|
return "5", "184"
|
|
if freq == "4925":
|
|
return "5", "185"
|
|
if freq == "4935":
|
|
return "5", "187"
|
|
if freq == "4940":
|
|
return "5", "188"
|
|
if freq == "4945":
|
|
return "5", "189"
|
|
if freq == "4960":
|
|
return "5", "192"
|
|
if freq == "4980":
|
|
return "5", "196"
|
|
return None
|