#! /usr/bin/env python
import sys
import dbus
import shlex
import subprocess
from signal import SIGTERM
from os import kill
from time import sleep


# dbus constants.
WPAS_DBUS_SERVICE = "fi.epitest.hostap.WPASupplicant"
WPAS_DBUS_INTERFACE = "fi.epitest.hostap.WPASupplicant"
WPAS_DBUS_OPATH = "/fi/epitest/hostap/WPASupplicant"

WPAS_DBUS_INTERFACES_INTERFACE = "fi.epitest.hostap.WPASupplicant.Interface"
WPAS_DBUS_INTERFACES_OPATH = "/fi/epitest/hostap/WPASupplicant/Interfaces"
WPAS_DBUS_BSSID_INTERFACE = "fi.epitest.hostap.WPASupplicant.BSSID"
WPAS_DBUS_NETWORKS_INTERFACE = "fi.epitest.hostap.WPASupplicant.Network"


class SupplicantStartError(Exception):
    pass

class WpaSupplicant:
    def __init__(self, interface, driver, extra_args=None):
        args=["wpa_supplicant", "-uB", "-P/var/run/wpa_supplicant.pid"] + extra_args

        # Start wpa_supplicant
        supplicant = subprocess.Popen(args,stderr=subprocess.STDOUT,stdout=subprocess.PIPE)
        output = supplicant.communicate()[0]
        if supplicant.returncode not in [255,0]:
            print output
            print " - Could not start wpa_supplicant"
            raise SupplicantStartError

        self.bus = dbus.SystemBus()
        wpas_obj = self.bus.get_object(WPAS_DBUS_SERVICE, WPAS_DBUS_OPATH)
        self.wpas = dbus.Interface(wpas_obj, WPAS_DBUS_INTERFACE)
        self.iface = self.get_iface(interface, driver)
        
    def get_iface(self, interface, driver):
        #iface=get_supplicant_iface(interface)
        try:
            path = self.wpas.getInterface(interface)
        except dbus.exceptions.DBusException:
            path = self.wpas.addInterface(interface, {"driver":dbus.String(driver,variant_level=1)})

        if_obj = self.bus.get_object(WPAS_DBUS_SERVICE, path)
        self.iface = dbus.Interface(if_obj, WPAS_DBUS_INTERFACES_INTERFACE);  

    def add_wpa_network(self, essid, key):
        path = iface.addNetwork()
        net_obj = bus.get_object(WPAS_DBUS_SERVICE, path)
        rnet = dbus.Interface(net_obj, WPAS_DBUS_NETWORKS_INTERFACE)
        iface.selectNetwork(rnet)

        opts = dbus.Dictionary({"ssid": dbus.ByteArray(essid), "psk": dbus.String(key)}, signature="sv")
        rnet.set(opts)

    def add_wep_network(self, essid, key):  
        path = iface.addNetwork()
        net_obj = bus.get_object(WPAS_DBUS_SERVICE, path)
        rnet = dbus.Interface(net_obj, WPAS_DBUS_NETWORKS_INTERFACE)
        iface.selectNetwork(rnet)

        if key[:2] == "s:": # String key
            key=key[2:]
        else: # Hex key
            key=wep_hex2dec(key)
        opts = dbus.Dictionary({"ssid": dbus.ByteArray(essid), "wpa_key0": dbus.ByteArray(key)}, signature="sv")
        rnet.set(opts)

    def get_status(self):
        return self.iface.status()


def read_config(config):
    cfg = shlex.split(open(config, "r").read())
    options = {}
    for line in cfg:
        (var, delim, value) = line.partition('=')  
        if delim and var.lstrip()[0] != "#":
            options[var] = value
    return options

def wep_hex2dec(key):
    if len(key) not in [10, 26]:
        print "Bad key"
        raise something
    
    x=0
    new_key=[]
    while x<len(key):
        new_key.append(int(key[x:x+2]),16)
        x+=2
    return new_key


def start(profile):
    # TODO: Add check if it's even a wireless interface
    # TODO: Do I need to bring the interface up first?

    if not profile ['SECURITY'] in ['wpa','wep','wpa-config']:
        print " - Invalid SECURITY specified. Valid: wpa, wep, wpa-config, none"
        return False
    
    args=[]
    try:
        args.append(profile['WPA_OPTS'])
    except KeyError:
        args.append("-Dwext")
    try:
        driver=profile["WPA_DRIVER"]
    except KeyError:
        driver="wext"
    if profile['SECURITY'] == "wpa-config":
        try:
            args.append("-c" + profile["WPA_CONF"])
        except KeyError:
            args.append("-c/etc/wpa_supplicant.conf")

    try:
        supplicant=WpaSupplicant(profile['INTERFACE'], driver, args)
    except SupplicantStartError:
        print " - WPA Supplicant did not start"
        return False

    if profile['SECURITY'] == "wep":
         supplicant.add_wpa_network(profile['ESSID'],profile['KEY'])
    elif profile['SECURITY'] == "wep":
        supplicant.add_wep_network(essid, key)

    # Determine timeout
    try:
        timeout = int(profile["TIMEOUT"])
    except KeyError:
        timeout = 15

    # Check for association
    for t in range(timeout):
        sleep(1)
        status = supplicant.get_status()
        if status == "COMPLETED":
            break
        elif t+1 == timeout:
            print " - Association/Auth failed:", status
            return False
            
    # Run ethernet and get an ip.
    try:
        subprocess.check_call(["/usr/lib/network/connections/ethernet-iproute", "up", sys.argv[2]])
    except subprocess.CalledProcessError:
        return False
    return True
  
def stop(profile):
    subprocess.call(["/usr/lib/network/connections/ethernet", "down", sys.argv[2]])
    kill(int(open("/var/run/wpa_supplicant.pid").read()),SIGTERM)
    return True 

if __name__ == "__main__":
    profile = read_config("/etc/network.d/"+sys.argv[2])
 
    if sys.argv[1] == "up": # Invert return codes for bash
        sys.exit(not start(profile))
    elif sys.argv[1] == "down":
        sys.exit(not stop(profile))
