#!/usr/bin/python
'''
//=============================================================================
//
//  File : ufo_wpasupplicant.py
//  Creation date : Thu Gen 02 14:30:48 CEST 2012
//  Working on this file:Grifisx (Antonino G. Imbesi)
//                      
//  This file is part of the Ufo Wardriving distribution
//
//  Websites: http://ufo-wardriving.com
//
//  This program is FREE software. You can redistribute it and/or
//  modify it under the terms of the GNU General Public License
//  as published by the Free Software Foundation; either version 2
//  of the License, or (at your opinion) any later version.
//
//  This program is distributed in the HOPE that it will be USEFUL,
//  but WITHOUT ANY WARRANTY; without even the implied warranty of
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
//  See the GNU General Public License for more details.
//
//  You should have received a copy of the GNU General Public License
//  along with this program. If not, write to the Free Software Foundation,
//  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
//
//=============================================================================
@ class name:
WPASupplicantObj:
@ short description:
Provides a class to handle the internet connection by wpa_supplicand
@ functions:
def setStatus(self,<str>):
    Set the status for current connection/class, emit statusChenged signal
        status for WPASupplicantObj class:
        -NEED_ROOT : initialization failed you must be root
        -NO_KEY : initialization failed you must have wifi key
        -INIT_OK  : all elements are ok but you are off line
        -WPA_COMPLETED: WAP is associated and 4whs is complete 
        -ON_LINE: you are on line
        -NO_ESSID: a no valid essid was given
        -DISCONNECTED: disconnected
        -DISABLE_NM: distro network manager was disabled
        -ENABLE_NM: distro network manager was enabled
        -IP_ERROR: no IP was given dns error?
-----------------------------------------------------------------
def printStatus(self,<str>):
    Print the current connection/class status  
-----------------------------------------------------------------
def <status> getStatus(self):
    Retrive the current connection/class status
-----------------------------------------------------------------
def <bool> isSuitable(self):
    Return True if all class elements are ok for connection
-----------------------------------------------------------------
def <bool> isRoot(self):
    Return True if running as root
-----------------------------------------------------------------
def <bool> checkInternetConnection(self):
    Check if internet connection is ok then return True
-----------------------------------------------------------------
def <bool> waitWpaConnection(self,<int = 5>):
    Wait for wpaconnection's "COMPLETED" status then return True
    default timeout is 5 sec
-----------------------------------------------------------------
def <WAP_STATUS> getWpaStatus(self):
    Retrive WPA Status, valid state are COMPLETED, ASSOCIATED
    ASSOCIATING, 4_WAY_HANDSHAKE etc.. etc..
-----------------------------------------------------------------
def <list> scanning(self):
-----------------------------------------------------------------
def makeWpaConfigFile(self):
-----------------------------------------------------------------
def makeWepConfigFile(self):
-----------------------------------------------------------------
def makeOpenConfigFile(self):
-----------------------------------------------------------------
def execWpaSupplicant(self, <str>):
-----------------------------------------------------------------
def connection(self):
-----------------------------------------------------------------
def <str(iface_name)> getIwface(self):
-----------------------------------------------------------------
def <str(iface_list)> getIwlist(self):
-----------------------------------------------------------------
def <iface_name> getPreferedIwface(self):
-----------------------------------------------------------------
def resetConnection(self,case=""):
-----------------------------------------------------------------
def <bool> startWpaSupplicant(self):
-----------------------------------------------------------------
def runDhClient(self):

@signals:
emit statusChanged <str(status)>:
    The signal is emitted when the object's state change in
    setStatus() function
'''

from PyQt4.QtCore import *
from PyQt4.QtGui import *
from core.ufo_define import *

import sys, os.path,platform, time
import subprocess
import urllib2

workingPath= str(QString.fromUtf8(repr(os.path.dirname(\
        os.path.realpath(sys.argv[0]))).replace("\\\\","/")\
        .replace("\'","")+"/config/"))
LOG_FILE = workingPath+"wpalog.log"
CFG_FILE = workingPath+"wpa.conf"

def dbg(arg):
    if DEFINE_DEBUG_MODE:
        print "@"+str(arg)
            
class WPASupplicantObj(QObject):
    def __init__(self, essid, key="", encr="Open", iface="", parent = None):
        super(WPASupplicantObj, self).__init__(parent=None)
        self.parent = parent
        self.bHaveKey = True
        if key =="":
            self.bHaveKey = False
        if self.isRoot() == False:
            self.setStatus("NEED_ROOT")
            return 
        self.essid = essid
        if self.bHaveKey == False and encr !="Open" :
            self.setStatus("NO_KEY")
            return
        if essid == " - " or essid == "":
            self.setStatus("NO_ESSID")
            return
        self.key=key
        self.encr=encr
        self.iface=iface
        self.key_mgmt =""
        if self.iface == "":
            self.iface = self.getPreferedIwface()
        self.wpastatus = ""
        self.connectionstatus = ""
        dbg(self.essid+" - "+self.key+" - "+self.encr+" - "+self.iface)
        self.makeConfigFile()
        if self.encr == "WPA" or self.encr == "WPA2":
            self.key_mgmt ="WPA-PSK"
        elif self.encr == "WEP":
            self.key_mgmt ="NONE"
        else:
            self.key_mgmt ="NONE"
            
        self.makeConfigFile()
        self.setStatus("INIT_OK")

    def setStatus(self,status):
        self.status = status
        self.emit(SIGNAL("statusChanged"), status)

    def getStatus(self):
        return self.status
        
    def isSuitable(self):
        if self.status == "INIT_OK" or self.status == "ON_LINE" or self.status == "WPA_COMPLETED":
            return True
        else:
            print "!! "+self.status+" !!"
            return False
        
    def isRoot(self):
        if os.getuid() is not 0:
            return 0
        else:
            return 1

    def checkInternetConnection(self):
        try:
            response=urllib2.urlopen('http://www.google.com',timeout=1)
            return True
        except urllib2.URLError as err: pass
        return False

    def wait_wpaConnection(self,timeout = 5):
        end_time = time.time() + timeout
        while time.time() < end_time:
            ws = self.getWpaStatus()
            if ws == None: return False
            if ws != "COMPLETED":
                self.setStatus(ws)
            else:
                self.setStatus("WPA_COMPLETED")
                return True
            time.sleep(1)
        return False

    def getWpaStatus(self):
        popen = subprocess.Popen(['wpa_cli','status'], stdout = subprocess.PIPE)
        popen.wait()
        (stdout,stderr) = popen.communicate(None)
        wpastatus = ""
        for line in stdout.split('\n'):
            key = line.split('=')[0]
            if key == 'wpa_state':
                wpastatus = line.split('=')[1]
                return wpastatus

    def scanning(self):
        popen = subprocess.Popen(['wpa_cli','scan_results'], stdout = subprocess.PIPE)
        popen.wait()
        (stdout,stderr) = popen.communicate(None)
        availablenets = stdout.split('\n')[2:]
        return availablenets

    def makeConfigFile(self):
        if self.fileExists(CFG_FILE): subprocess.call(["rm",CFG_FILE])
        f = QFile(CFG_FILE)
        if not f.open(QIODevice.ReadWrite) : dbg("Failed!")
        t = QTextStream(f)
        t << "ctrl_interface=/var/run/wpa_supplicant \n\n"
        t << "network={ \n"
        t << " ssid=\""+self.essid+"\"\n"
        if self.encr == "WPA" or self.encr == "WPA2":
            t << " psk=\""+self.key+"\"\n"
        if self.encr != "Open":
            t << " proto="+self.encr+"\n"
        t << " key_mgmt="+self.key_mgmt+"\n"
        if self.encr == "WEP":
            t << " wep_key0="+self.key
        t << "}"
        f.close();

    def connection(self):
        if not self.isSuitable():
            return
        if not self.iface:
            self.iface = self.getPreferedIwface()
        self.resetConnection("INIT_CONNECTION")
        wc = self.startWpaSupplicant()
        if wc == False:
            self.setStatus("NO_WPA_CONNECTION")
            return False 
        else:
            self.setStatus("WPA_COMPLETED")
            ic = self.checkInternetConnection()
            if ic == False:
                self.runDhClient()
                icc = self.checkInternetConnection()
                if icc == False:
                    self.setStatus("IP_ERROR")
                    return False
                else:
                    self.setStatus("ON_LINE")
                    return True
            else:
                self.setStatus("ON_LINE")
                return True

    def getIwface(self):   
        dbg("WARNING! UNTRUST FUNCTION!")
        scan_net = subprocess.Popen("iwconfig", stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell = True)
        schede_rete = scan_net.stdout.readlines()
        schede_rete = [x.split(None, 1)[0] for x in schede_rete if 'IEEE 802.11' in x]
        for scheda in schede_rete:
            scheda = scheda[0:]
        return scheda

    def getIwlist(self):
        p = subprocess.Popen("wpa_cli interface", stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell = True)
        iwfaceList = p.stdout.readlines()
        return iwfaceList[2:]

    def getPreferedIwface(self):
        p = subprocess.Popen("wpa_cli interface", stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell = True)
        p.wait()
        iface = self.iface
        (stdout,stderr) = p.communicate(None)
        
        for line in stdout.split('\n'):
            if line.startswith('Selected interface'):
                self.iface = line.split("'")[1]
        if self.iface =="": self.iface = self.getIwface()
        return self.iface

    def resetConnection(self,case=""):
        self.reset_iface()
        
        dbg("Killing wpa_supplicant..")
        subprocess.call(["wpa_cli","terminate"])
        subprocess.call(["killall","wpa_supplicant"])

        dbg("Delete wpa_supplicant files..")
        if self.fileExists(LOG_FILE): subprocess.call(["rm",LOG_FILE])

        ifc ='/var/run/wpa_supplicant/'+self.iface
        if self.fileExists(ifc): subprocess.call(["rm",ifc])
        
        if case != "INIT_CONNECTION":
            if self.fileExists(CFG_FILE): subprocess.call(["rm",CFG_FILE])
        
        subprocess.call(["killall", "dhclient"])
        
        if case == "INIT_CONNECTION":
           self.setStatus("INIT_OK")
        else:
            self.setStatus("DISCONNECTED")

    def reset_iface(self):
        subprocess.call(["ifconfig",self.iface,"down"])
        subprocess.call(["ifconfig",self.iface,"up"])
        
    def enable_network_manager(self):
        if DEFINE_DISTRO_FEDORA == False:
            subprocess.call(["start","network-manager"])
        else:
            subprocess.call(["/etc/init.d/NetworkManager","start"])
        self.setStatus("ENABLE_NM")

    def disable_network_manager(self):
        if DEFINE_DISTRO_FEDORA == False:
            subprocess.call(["stop","network-manager"])
        else:
            subprocess.call(["/etc/init.d/NetworkManager","stop"])
        self.setStatus("DISABLE_NM")

    def startWpaSupplicant(self):
        cmd='wpa_supplicant -c "'+CFG_FILE+'" -D wext -dd -i '+self.iface+' -B > '+LOG_FILE
        dbg(str("Start wpa_supplicant"+cmd))
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell = True)
        return_code = proc.wait()
        dbg("WPA said: --> "+str(return_code))
        time.sleep(2)
        wc = self.wait_wpaConnection()
        if wc == False:
            return False 
        else:
            return True

    def runDhClient(self):
        dbg("Running dhcp...")
        subprocess.check_call(["dhclient", self.iface])
        
    def fileExists(self,filename):
        try:
           with open(filename) as f: return True
        except IOError as e:
           dbg("File Not Found: "+filename)
           return False

