From 4900be956057415c9597ff752b1781fad18e959c Mon Sep 17 00:00:00 2001
From: Delirious Lettuce <delirious.lettuce@gmail.com>
Date: Mon, 26 Jun 2017 19:40:54 -0600
Subject: [PATCH] Various fixes

---
 scripts/rpi_network_conn.py | 112 +++++++++++++++++++-----------------
 1 file changed, 59 insertions(+), 53 deletions(-)

diff --git a/scripts/rpi_network_conn.py b/scripts/rpi_network_conn.py
index 6c3d4b2..5d44b77 100755
--- a/scripts/rpi_network_conn.py
+++ b/scripts/rpi_network_conn.py
@@ -1,94 +1,100 @@
+import textwrap
 import time
+import urllib2
 from subprocess import CalledProcessError, check_output, Popen
 
-import urllib2
 import wifi_captive_portal
 
 
 def get_wifi_list():
+    # check wlan status, return [] if no wifi available
     try:
-        wlan_status = check_output(["sudo", "ifup", "wlan0"])
-        returncode = 0
-    except CalledProcessError as e:
-        returncode = e.returncode
+        check_output(["sudo", "ifup", "wlan0"])
+    except CalledProcessError:
+        return []
 
+    # check for wifi access points in range, return [] if none available
     try:
-        iw_list = (check_output(["sudo", "iwlist", "wlan0", "scan"])).split("\n")
-    except CalledProcessError as e:
-        iw_list = []
+        iw_list = check_output(["sudo", "iwlist", "wlan0", "scan"]).split("\n")
+    except CalledProcessError:
+        return []
 
     # contains a tuple of the (ESSID, Encryption key)
     wifi_list = []
 
-    for x in range(0, len(iw_list)):
-        if (iw_list[x].strip())[0:5] == "ESSID":
-            if (iw_list[x].strip())[7:-1] != "" and (not (iw_list[x].strip())[7:-1].startswith("\\")):
-                wifi_list.append((((iw_list[x].strip())[7:-1] + "-" + ((iw_list[x-1].strip())[15:])), (iw_list[x].strip())[7:-1]))
+    for i, interface_info in enumerate(iw_list):
+        current = interface_info.strip()
+        if current.startswith("ESSID"):
+            end = current[7:-1]
+            if end and (not end.startswith("\\")):
+                wifi_list.append(
+                    ('%s-%s' % (end, iw_list[i - 1].strip()[15:]), end))
 
-    return sorted(list(set(wifi_list)), key=lambda wifilist: wifilist[0])
+    return sorted(set(wifi_list), key=lambda wifilist: wifilist[0])
 
 
 def add_wifi(wifi_ssid, wifi_key):
+    ssid = 'ssid="%s"' % wifi_ssid
+
+    if wifi_key == "key_mgmt_none":
+        wifi_key_info = "    key_mgmt=NONE\n"
+    else:
+        wifi_key_info = '    psk="%s"\n' % wifi_key
+
     with open("/etc/wpa_supplicant/wpa_supplicant.conf") as wpa_supplicant:
         lines = wpa_supplicant.readlines()
-    
+
     wifi_exists = False
-    for x in range(0, len(lines)):
-        #if SSID is already in file
-        if (lines[x].strip()) == 'ssid="' + str(wifi_ssid) + '"':
+    for i in xrange(len(lines)):
+        # if SSID is already in file
+        if lines[i].strip() == ssid:
             wifi_exists = True
-            lines[x] = '    ssid="%s"\n' % wifi_ssid
-            if wifi_key == "key_mgmt_none":
-                lines[x+1] = '    key_mgmt=NONE\n'
-            else:
-                lines[x+1] = '    psk="%s"\n' % wifi_key
-
-    if wifi_exists is not True:
-
-        lines.append('network={\n')
-        lines.append('    ssid="%s"\n' % wifi_ssid)
-        if wifi_key == "key_mgmt_none":
-            lines.append('    key_mgmt=NONE\n')
-        else:
-            lines.append('    psk="%s"\n' % wifi_key)
-        lines.append('}\n')
+            lines[i] = "    %s\n" % ssid
+            lines[i + 1] = wifi_key_info
+
+    if not wifi_exists:
+        lines.append(textwrap.dedent("""\
+            network={
+                %s
+                %s
+            }
+            """ % (ssid, wifi_key_info)))
 
     with open("/etc/wpa_supplicant/wpa_supplicant.conf", "w") as fout:
-        for line in lines:
-            fout.write(line)
-    
+        fout.writelines(lines)
+
     process = Popen(["sudo", "ifdown", "wlan0"])
     process.wait()
     process = Popen(["sudo", "ifup", "wlan0"])
     process.wait()
-    
+
     time.sleep(15)
-    
+
     if not internet_status():
         wifi_captive_portal.captive_portal(wifi_ssid, "", "")
 
 
 def internet_status():
     try:
-        response = urllib2.urlopen("https://aws.amazon.com", timeout=1)
+        urllib2.urlopen("https://aws.amazon.com", timeout=1)
+    except urllib2.URLError:
+        return False
+    else:
         return True
-    except urllib2.URLError as err:
-        pass
-    return False
 
 
 def reset_wifi():
-    lines = ["country=US\n",
-             "ctrl_interface=DIR=/var/run/wpa_supplicant GROUP=netdev\n",
-             "update_config=1\n"]
+    lines = textwrap.dedent("""\
+        country=US
+        ctrl_interface=DIR=/var/run/wpa_supplicant GROUP=netdev
+        update_config=1
+        """)
     with open("/etc/wpa_supplicant/wpa_supplicant.conf", "w") as fout:
-        for line in lines:
-            fout.write(line)
-    
-    try:
-        vpn_status = (check_output(["sudo", "ifdown", "wlan0"]))
-        returncode = 0
-    except CalledProcessError as e:
-        returncode = e.returncode
+        fout.write(lines)
 
-    return returncode == 0
+    try:
+        check_output(["sudo", "ifdown", "wlan0"])
+    except CalledProcessError:
+        return False
+    else:
+        return True