#!/usr/bin/python

from py.Configuration import Configuration
from py.Scanner import Scanner
from py.Color import Color
from py.AttackWEP import AttackWEP
from py.AttackWPA import AttackWPA
from py.AttackWPS import AttackWPS
from py.CrackResult import CrackResult
from py.Handshake import Handshake

from json import loads
import os

class Wifite(object):

    def main(self):
        ''' Either performs action based on arguments, or starts attack scanning '''
        Configuration.initialize(load_interface=False)

        if Configuration.show_cracked:
            self.display_cracked()

        elif Configuration.check_handshake:
            self.check_handshake(Configuration.check_handshake)

        elif Configuration.crack_wpa:
            # TODO: Crack .cap file at crack_wpa
            Color.pl('{!} Unimplemented method: crack_wpa')
            pass
        elif Configuration.crack_wep:
            # TODO: Crack .cap file at crack_wep
            Color.pl('{!} Unimplemented method: crack_wep')
            pass
        elif Configuration.update:
            # TODO: Get latest version from github
            Color.pl('{!} Unimplemented method: update')
            pass
        else:
            self.run()

    def display_cracked(self):
        ''' Show cracked targets from cracked.txt '''
        Color.pl('{+} displaying {C}cracked target(s){W}')
        name = CrackResult.cracked_file
        if not os.path.exists(name):
            Color.pl('{!} {O}file {C}%s{O} not found{W}' % name)
            return
        f = open(name, 'r')
        json = loads(f.read())
        f.close()
        for (index, item) in enumerate(json):
            Color.pl('\n{+} Cracked target #%d:' % (index + 1))
            cr = CrackResult.load(item)
            cr.dump()
        
    def check_handshake(self, capfile):
        ''' Analyzes .cap file for handshake '''
        Color.pl('{+} checking for handshake in .cap file {C}%s{W}' % capfile)
        if not os.path.exists(capfile):
            Color.pl('{!} {O}.cap file {C}%s{O} not found{W}' % capfile)
            return
        hs = Handshake(capfile)
        hs.analyze()


    def run(self):
        '''
            Main program.
            1) Scans for targets, asks user to select targets
            2) Attacks each target
        '''
        s = Scanner()
        if s.target:
            # We found the target we want
            targets = [s.target]
        else:
            targets = s.select_targets()

        for t in targets:
            Color.pl('{+} starting attacks against {C}%s{W} ({C}%s{W})'
                % (t.bssid, t.essid))
            if 'WEP' in t.encryption:
                attack = AttackWEP(t)
                attack.run()
            elif 'WPA' in t.encryption:
                if t.wps:
                    attack = AttackWPS(t)
                    if attack.run():
                        # We cracked it.
                        break
                    else:
                        # WPS failed, try WPA handshake.
                        attack = AttackWPA(t)
                        attack.run()
                else:
                    # Not using WPS, try WPA handshake.
                    attack = AttackWPA(t)
                    attack.run()
            else:
                Color.pl("{!} {R}Error: {O}unable to attack: encryption not WEP or WPA")
                continue
                pass

            if attack.success:
                attack.crack_result.save()
        pass

    def print_banner(self):
        """ Displays ASCII art of the highest caliber.  """
        Color.pl("")
        Color.pl("{G}  .;'                     `;,    ")
        Color.pl("{G} .;'  ,;'             `;,  `;,  " +
            "{W}WiFite v%.2f" % Configuration.version)
        Color.pl("{G}.;'  ,;'  ,;'     `;,  `;,  `;,  ")
        Color.pl("{G}::   ::   :   {GR}( ){G}   :   ::   ::  " +
            "{W}automated wireless auditor")
        Color.pl("{G}':.  ':.  ':. {GR}/_\\{G} ,:'  ,:'  ,:'  ")
        Color.pl("{G} ':.  ':.    {GR}/___\\{G}   ,:'  ,:'   " +
            "{W}designed for Linux")
        Color.pl("{G}  ':.       {GR}/_____\\{G}     ,:'     ")
        Color.pl("{G}           {GR}/       \\{G}            ")
        Color.pl("{W}")


if __name__ == '__main__':
    w = Wifite()
    try:
        w.print_banner()
        w.main()
    except Exception, e:
        Color.pl('\n{!} {R}Error:{O} %s{W}' % str(e))
        from traceback import format_exc
        print '\n    '
        print format_exc().replace('\n', '\n    ')
    except KeyboardInterrupt:
        Color.pl('\n{!} {O}interrupted{W}')
    Configuration.exit_gracefully(0)