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

import logging
import socket
from os.path import isfile
from threading import BoundedSemaphore
from socket import gethostbyname
from functools import wraps
from cme.logger import CMEAdapter
from cme.context import Context
from cme.helpers.logger import write_log

sem = BoundedSemaphore(1)
global_failed_logins = 0
user_failed_logins = {}


def gethost_addrinfo(hostname):
    try:
        for res in socket.getaddrinfo(hostname, None, socket.AF_INET6,
                socket.SOCK_DGRAM, socket.IPPROTO_IP, socket.AI_CANONNAME):
            af, socktype, proto, canonname, sa = res
    except socket.gaierror:
        for res in socket.getaddrinfo(hostname, None, socket.AF_INET,
                socket.SOCK_DGRAM, socket.IPPROTO_IP, socket.AI_CANONNAME):
            af, socktype, proto, canonname, sa = res
    if canonname == '':
        return sa[0]
    return canonname


def requires_admin(func):
    def _decorator(self, *args, **kwargs):
        if self.admin_privs is False:
            return
        return func(self, *args, **kwargs)
    return wraps(func)(_decorator)


class connection(object):

    def __init__(self, args, db, host):
        self.domain = None
        self.args = args
        self.db = db
        self.hostname = host
        self.conn = None
        self.admin_privs = False
        self.logger = None
        self.password = ''
        self.username = ''
        self.kerberos = True if self.args.kerberos or self.args.use_kcache else False
        self.aesKey = None if not self.args.aesKey else self.args.aesKey
        self.kdcHost = None if not self.args.kdcHost else self.args.kdcHost
        self.log = None if not self.args.log else self.args.log
        self.use_kcache = None if not self.args.use_kcache else self.args.use_kcache
        self.failed_logins = 0
        self.local_ip = None

        if self.log:
            CMEAdapter().setup_logfile(self.log[0])

        try:
            self.host = gethost_addrinfo(self.hostname)
            if self.args.kerberos:
                self.host = self.hostname
        except Exception as e:
            logging.debug('Error resolving hostname {}: {}'.format(self.hostname, e))
            return

        self.proto_flow()

    @staticmethod
    def proto_args(std_parser, module_parser):
        return

    def proto_logger(self):
        pass

    def enum_host_info(self):
        return

    def print_host_info(info):
        return

    def create_conn_obj(self):
        return

    def check_if_admin(self):
        return

    def kerberos_login(self):
        return

    def plaintext_login(self, domain, username, password):
        return

    def hash_login(self, domain, username, ntlm_hash):
        return

    def proto_flow(self):
        if self.create_conn_obj():
            self.enum_host_info()
            self.proto_logger()
            if self.print_host_info():
                # because of null session
                if self.login() or (self.username == '' and self.password == ''):
                    if hasattr(self.args, 'module') and self.args.module:
                        self.call_modules()
                    else:
                        self.call_cmd_args()

    def call_cmd_args(self):
        for k, v in vars(self.args).items():
            if hasattr(self, k) and hasattr(getattr(self, k), '__call__'):
                if v is not False and v is not None:
                    logging.debug('Calling {}()'.format(k))
                    r = getattr(self, k)()

    def call_modules(self):
        module_logger = CMEAdapter(extra={
            'module': self.module.name.upper(),
            'host': self.host,
            'port': self.args.port,
            'hostname': self.hostname
        })

        context = Context(self.db, module_logger, self.args)
        context.localip = self.local_ip

        if hasattr(self.module, 'on_request') or hasattr(self.module, 'has_response'):
            self.server.connection = self
            self.server.context.localip = self.local_ip

        if hasattr(self.module, 'on_login'):
            self.module.on_login(context, self)

        if self.admin_privs and hasattr(self.module, 'on_admin_login'):
            self.module.on_admin_login(context, self)

        if (not hasattr(self.module, 'on_request') and not hasattr(self.module, 'has_response')) and hasattr(self.module, 'on_shutdown'):
            self.module.on_shutdown(context, self)

    def inc_failed_login(self, username):
        global global_failed_logins
        global user_failed_logins

        if username not in user_failed_logins.keys():
            user_failed_logins[username] = 0

        user_failed_logins[username] += 1
        global_failed_logins += 1
        self.failed_logins += 1

    def over_fail_limit(self, username):
        global global_failed_logins
        global user_failed_logins

        if global_failed_logins == self.args.gfail_limit:
            return True

        if self.failed_logins == self.args.fail_limit:
            return True

        if username in user_failed_logins.keys():
            if self.args.ufail_limit == user_failed_logins[username]:
                return True

        return False

    def login(self):
        for cred_id in self.args.cred_id:
            with sem:
                if cred_id.lower() == 'all':
                    creds = self.db.get_credentials()
                else:
                    creds = self.db.get_credentials(filter_term=int(cred_id))
                for cred in creds:
                    logging.debug(cred)
                    try:
                        c_id, domain, username, password, cred_type, pillaged_from = cred

                        if cred_type and password:
                            if not domain:
                                domain = self.domain

                            if self.args.local_auth:
                                domain = self.domain
                            elif self.args.domain:
                                domain = self.args.domain

                            if cred_type == 'hash' and not self.over_fail_limit(username):
                                if self.args.kerberos:
                                    if self.kerberos_login(domain, username, '', password, '', self.kdcHost, False):
                                        return True
                                elif self.hash_login(domain, username, password): return True

                            elif cred_type == 'plaintext' and not self.over_fail_limit(username):
                                if self.args.kerberos:
                                    if self.kerberos_login(domain, username, password, '' , '', self.kdcHost, False):
                                        return True
                                elif self.plaintext_login(domain, username, password): return True
                    except IndexError:
                        self.logger.error("Invalid database credential ID!")
        if self.args.use_kcache:
            with sem:
                username = self.args.username[0] if len(self.args.username) else ''
                password = self.args.password[0] if len(self.args.password) else ''
                self.kerberos_login(self.domain, username, password, '', '', self.kdcHost, True)
                return True
        for user in self.args.username:
            if isfile(user):
                with open(user, 'r') as user_file:
                    for usr in user_file:
                        if "\\" in usr:
                            tmp = usr
                            usr = tmp.split('\\')[1].strip()
                            self.domain = tmp.split('\\')[0]
                        if hasattr(self.args, 'hash') and self.args.hash:
                            with sem:
                                for ntlm_hash in self.args.hash:
                                    if isfile(ntlm_hash):
                                        with open(ntlm_hash, 'r') as ntlm_hash_file:
                                            if not self.args.no_bruteforce:
                                                for f_hash in ntlm_hash_file:
                                                    if not self.over_fail_limit(usr.strip()):
                                                        if self.args.kerberos:
                                                            if self.kerberos_login(self.domain, usr.strip(), '', f_hash.strip(), '', self.kdcHost, False):
                                                                return True
                                                        elif self.hash_login(self.domain, usr.strip(), f_hash.strip()):
                                                            return True
                                            elif self.args.no_bruteforce:
                                                user_file.seek(0) # HACK: this should really not be in the usr for loop
                                                for usr, f_hash in zip(user_file, ntlm_hash_file):
                                                    if not self.over_fail_limit(usr.strip()):
                                                        if self.args.kerberos:
                                                            if self.kerberos_login(self.domain, usr.strip(), '', f_hash.strip(), '', self.kdcHost, False): return True
                                                        elif self.hash_login(self.domain, usr.strip(), f_hash.strip()):
                                                            return True
                                    else: # ntlm_hash is a string
                                        if not self.over_fail_limit(usr.strip()):
                                            if self.args.kerberos:
                                                if self.kerberos_login(self.domain, usr.strip(), '', ntlm_hash.strip(), '', self.kdcHost, False):
                                                    return True
                                            elif self.hash_login(self.domain, usr.strip(), ntlm_hash.strip()):
                                                return True
                        elif self.args.password:
                            with sem:
                                for password in self.args.password:
                                    if isfile(password):
                                        with open(password, 'r') as password_file:
                                            if not self.args.no_bruteforce:
                                                for f_pass in password_file:
                                                    if not self.over_fail_limit(usr.strip()):
                                                        if hasattr(self.args, 'domain'):
                                                            if self.args.kerberos:
                                                                if self.kerberos_login(self.domain, usr.strip(), f_pass.strip(), '', '', self.kdcHost, False):
                                                                    return True
                                                            elif self.plaintext_login(self.domain, usr.strip(), f_pass.strip()):
                                                                return True
                                                        else:
                                                            if self.plaintext_login(usr.strip(), f_pass.strip()):
                                                                return True
                                            elif self.args.no_bruteforce:
                                                user_file.seek(0)  # HACK: this should really not be in the usr for loop
                                                for usr, f_pass in zip(user_file, password_file):
                                                    if not self.over_fail_limit(usr.strip()):
                                                        if hasattr(self.args, 'domain'):
                                                            if self.args.kerberos:
                                                                if self.kerberos_login(self.domain, usr.strip(), f_pass.strip(), '', '', self.kdcHost, False):
                                                                    return True
                                                            elif self.plaintext_login(self.domain, usr.strip(), f_pass.strip()):
                                                                return True
                                                        else:
                                                            if self.plaintext_login(usr.strip(), f_pass.strip()):
                                                                return True
                                    else:  # password is a string
                                        if not self.over_fail_limit(usr.strip()):
                                            if hasattr(self.args, 'domain'):
                                                if self.args.kerberos:
                                                    if self.kerberos_login(self.domain, usr.strip(), password, '', '', self.kdcHost, False):
                                                        return True
                                                elif self.plaintext_login(self.domain, usr.strip(), password):
                                                    return True
                                            else:
                                                if self.plaintext_login(usr.strip(), password):
                                                    return True
            else:  # user is a string
                if hasattr(self.args, 'hash') and self.args.hash:
                    with sem:
                        for ntlm_hash in self.args.hash:
                            if isfile(ntlm_hash):
                                with open(ntlm_hash, 'r') as ntlm_hash_file:
                                    for f_hash in ntlm_hash_file:
                                        if not self.over_fail_limit(user):
                                            if self.args.kerberos:
                                                if self.kerberos_login(self.domain, user, '', ntlm_hash.strip(), '', self.kdcHost, False):
                                                    return True
                                            elif self.hash_login(self.domain, user, f_hash.strip()):
                                                return True
                            else: # ntlm_hash is a string
                                if not self.over_fail_limit(user):
                                    if self.args.kerberos:
                                        if self.kerberos_login(self.domain, user, '', ntlm_hash.strip(), '', self.kdcHost, False): return True
                                    elif self.hash_login(self.domain, user, ntlm_hash.strip()):
                                        return True
                elif self.args.password:
                    with sem:
                        for password in self.args.password:
                            if isfile(password):
                                with open(password, 'r') as password_file:
                                    for f_pass in password_file:
                                        if not self.over_fail_limit(user):
                                            if hasattr(self.args, 'domain'):
                                                if self.args.kerberos:
                                                    if self.kerberos_login(self.domain, user, f_pass.strip(), '', '', self.kdcHost, False):
                                                        return True
                                                elif self.plaintext_login(self.domain, user, f_pass.strip()):
                                                    return True
                                            else:
                                                if self.plaintext_login(user, f_pass.strip()):
                                                    return True
                            else:  # password is a string
                                if not self.over_fail_limit(user):
                                    if hasattr(self.args, 'domain'):
                                        if self.args.kerberos:
                                            if self.kerberos_login(self.domain, user, password, '', '', self.kdcHost, False):
                                                return True
                                        elif self.plaintext_login(self.domain, user, password):
                                            return True
                                    else:
                                        if self.plaintext_login(user, password):
                                            return True
                elif self.args.aesKey:
                    with sem:
                        for aesKey in self.args.aesKey:
                            if not self.over_fail_limit(user):
                                if self.kerberos_login(self.domain, user, '', '', aesKey.strip(), self.kdcHost, False):
                                    return True