figure out P"" strings for wpa_supplicant.conf exactly

diff --git a/select_wifi.py b/select_wifi.py
index d0c3a53cdadd6962a0231ccdb623ed67edb3d63a..28c51ec91522fd2baeae67a1d338902ea9be6255 100755 (executable)
--- a/select_wifi.py
+++ b/select_wifi.py
@@ -7,59 +7,107 @@ from subprocess import check_call
 WPA_CONFIG = "/etc/wpa_supplicant/wpa_supplicant.conf"
 
 
+def unescape(s):
+    hexdig = "0123456789abcdef"
+    octdig = hexdig[:8]
+    octesc = tuple(f"\\{c}" for c in octdig)
+    simple_map = {
+        '\\"': ("", 1),
+        "\\\\": ("", 1),
+        "\\n": ("\n", 2),
+        "\\r": ("\r", 2),
+        "\\t": ("\t", 2),
+        "\\e": ("\033", 2),
+    }
+    pos = None
+    while (pos := s.find("\\", pos)) != -1:
+        e = s[pos:pos + 2]
+        ch, rm = simple_map.get(e, ("", 2))
+        if e == "\\x":
+            hi = hexdig.find(s[pos + 2].lower())
+            ch = ""
+            if hi >= 0:
+                lo = hexdig.find(s[pos + 3].lower())
+                rm += 1
+                if lo >= 0:
+                    rm += 1
+                    ch = chr((hi << 4) | lo)
+                else:
+                    ch = chr(hi)
+        elif e in octesc:
+            chv = int(s[pos + 1])
+            if s[pos + 2] in octdig:
+                rm += 1
+                chv = 8 * chv + int(s[pos + 2])
+                if s[pos + 3] in octdig:
+                    rm += 1
+                    chv = 8 * chv + int(s[pos + 3])
+            ch = chr(chv)
+        s = f"{s[:pos]}{ch}{s[pos + rm:]}"
+        pos += 1
+    return s
+
+
 class WpaConfig:
     SSID_PREFIX = ("    #ssid=", "    ssid=")
     PSK_PREFIX = ("    #psk=", "    psk=")
     KEY_MGMT_PREFIX = ("    #key_mgmt=", "    key_mgmt=")
+    COMMENTED_PREFIX = (SSID_PREFIX[0], PSK_PREFIX[0], KEY_MGMT_PREFIX[0])
 
     def __init__(self):
         self.lines = []
 
-    class Line:
-        def __init__(self, s):
-            self.s = s
-
     @staticmethod
     def parse_str(s):
         if s.endswith('"'):
             if s.startswith('"'):
                 return s[1:-1]
             elif s.startswith('P"'):
-                return s[2:-1].encode().decode("unicode_escape")
+                return unescape(s[2:-1])
 
     def print(self, fh=None):
         if fh is None:
             fh = sys.stdout
         for line in self.lines:
-            print(line.s, file=fh)
+            print(line, file=fh)
 
     def choose(self):
         choices = []
         for i, line in enumerate(self.lines):
-            if not any(line.s.startswith(pfx) for pfx in self.SSID_PREFIX):
+            if not any(line.startswith(pfx) for pfx in self.SSID_PREFIX):
+                pos = line.find("#")
+                if pos != -1 and (
+                    pos != 4
+                    or not any(
+                        line.startswith(pfx) for pfx in self.COMMENTED_PREFIX
+                    )
+                ):
+                    print(f"     {line[pos:]}")
                 continue
-            ssid = self.parse_str(line.s[line.s.find("=") + 1:])
-            chosen = " " if line.s[4] == "#" else "*"
+            ssid = self.parse_str(line[line.find("=") + 1:])
+            chosen = " " if line[4] == "#" else "*"
             print(f"{len(choices):3} {chosen}{ssid}")
             choices.append((i, line))
         index, line = choices[int(input("> "))]
-        if line.s.startswith(self.SSID_PREFIX[1]):
-            return
-        for line in self.lines:
-            if line.s.startswith("    ") and line.s[4] != "#":
-                line.s = f"    #{line.s[4:]}"
-        self.lines[index].s = f"    {self.lines[index].s[5:]}"
+        if line.startswith(self.SSID_PREFIX[1]):
+            return False
+        for i, line in enumerate(self.lines):
+            if line.startswith("    ") and line[4] != "#":
+                self.lines[i] = f"    #{line[4:]}"
+        self.lines[index] = f"    {self.lines[index][5:]}"
         index += 1
-        if index < len(self.lines) and self.lines[index].s.startswith(
+        if index < len(self.lines) and self.lines[index].startswith(
             self.PSK_PREFIX[0]
         ):
-            self.lines[index].s = f"    {self.lines[index].s[5:]}"
-            return
-        for line in self.lines:
-            if line.s.startswith(self.KEY_MGMT_PREFIX[0]):
-                assert line.s[5:] == "key_mgmt=NONE"
-                line.s = f"    {line.s[5:]}"
-                return
+            self.lines[index] = f"    {self.lines[index][5:]}"
+            return True
+        for i, line in enumerate(self.lines):
+            if line.startswith(self.KEY_MGMT_PREFIX[0]):
+                assert line[5:] == "key_mgmt=NONE"
+                self.lines[i] = f"    {line[5:]}"
+                return True
+        self.lines.append("    key_mgmt=NONE")
+        return True
 
 
 def main():
@@ -67,8 +115,9 @@ def main():
     wpa_config = WpaConfig()
     with open(WPA_CONFIG) as fh:
         for line in fh:
-            wpa_config.lines.append(WpaConfig.Line(line.rstrip()))
-    wpa_config.choose()
+            wpa_config.lines.append(line.rstrip())
+    if not wpa_config.choose():
+        return
     with open(WPA_CONFIG, "wt") as fh:
         wpa_config.print(fh)
     check_call(["rc-service", "wpa_supplicant", "restart"])