#!/usr/bin/env python

import sys, os
import base64
import time
import select, errno
if sys.version_info < (3,):
    import urllib
    from urlparse import urlparse
else:
    import urllib.request, urllib.parse, urllib.error
    from urllib.parse import urlparse
import argparse
import socket
import subprocess
import syslog
from struct import unpack
from threading import Lock
from binascii import hexlify

from smb.smb_constants import *
from smb.smb_structs import *
from smb.smb2_structs import *
from smb.base import SMB, NotConnectedError, NotReadyError, SMBTimeout
from smb import ntlm, securityblob, smb_structs
import codecs

class SYNOIWAConnection(SMB):
    def __init__(self):
        SMB.__init__(self, '', '', '', '', '', True, SMB.SIGN_WHEN_REQUIRED, True)
        self.sock = None
        self.is_busy = False

    def onAuthOK(self):
        '''
        copy from SMBConnection
        '''
        self.auth_result = True

    def onAuthFailed(self):
        '''
        copy from SMBConnection
        '''
        self.auth_result = False

    def write(self, data):
        '''
        copy from SMBConnection
        '''
        assert self.sock
        data_len = len(data)
        total_sent = 0
        while total_sent < data_len:
            sent = self.sock.send(data[total_sent:])
            if sent == 0:
                raise NotConnectedError('Server disconnected')
            total_sent = total_sent + sent

    def close(self):
        '''
        copy from SMBConnection
        '''
        if self.sock:
            self.sock.close()
            self.sock = None

    def _pollForNetBIOSPacket(self, timeout):
        '''
        copy from SMBConnection
        '''
        expiry_time = time.time() + timeout
        read_len = 4
        data = b''

        while read_len > 0:
            try:
                if expiry_time < time.time():
                    raise SMBTimeout

                ready, _, _ = select.select([ self.sock.fileno() ], [ ], [ ], timeout)
                if not ready:
                    raise SMBTimeout

                d = self.sock.recv(read_len)
                if len(d) == 0:
                    raise NotConnectedError

                data = data + d
                read_len -= len(d)
            except select.error as ex:
                if isinstance(ex, types.TupleType):
                    if ex[0] != errno.EINTR and ex[0] != errno.EAGAIN:
                        raise ex
                else:
                    raise ex

        type_, flags, length = unpack('>BBH', data)
        if flags & 0x01:
            length = length | 0x10000

        read_len = length
        while read_len > 0:
            try:
                if expiry_time < time.time():
                    raise SMBTimeout

                ready, _, _ = select.select([ self.sock.fileno() ], [ ], [ ], timeout)
                if not ready:
                    raise SMBTimeout

                d = self.sock.recv(read_len)
                if len(d) == 0:
                    raise NotConnectedError

                data = data + d
                read_len -= len(d)
            except select.error as ex:
                if isinstance(ex, types.TupleType):
                    if ex[0] != errno.EINTR and ex[0] != errno.EAGAIN:
                        raise ex
                else:
                    raise ex

        self.feedData(data)

    def connect(self, ip, port = 445, sock_family = socket.AF_INET, timeout = 60):
        if self.sock:
            self.sock.close()

        self.auth_result = None
        self.sock = socket.socket(sock_family)
        self.sock.settimeout(timeout)
        self.sock.connect(( ip, port ))
        self.is_busy = True

        return self.sock

    def negotiate(self, ntlm_data, timeout = 60):
        self.ntlm_data = ntlm_data
        self.onNMBSessionOK()
        self._pollForNetBIOSPacket(timeout)
        self._pollForNetBIOSPacket(timeout)
        result, ntlm_token = securityblob.decodeChallengeSecurityBlob(self.smb_message.payload.security_blob)

        return ntlm_token

    def authenticate(self, ntlm_data, timeout = 60):
        if smb_structs.SUPPORT_SMB2:
            self.authenticate_SMB2(ntlm_data)
        else:
            self.authenticate_SMB1(ntlm_data)
        self._pollForNetBIOSPacket(timeout)
        self.is_busy = False

        return self.auth_result

    def authenticate_SMB1(self, ntlm_data):
        assert self.smb_message.hasExtendedSecurity

        if self.smb_message.uid and not self.uid:
            self.uid = self.smb_message.uid

        blob = securityblob.generateAuthSecurityBlob(ntlm_data)
        self._sendSMBMessage(SMBMessage(ComSessionSetupAndxRequest__WithSecurityExtension(0, blob)))

        if self.security_mode & NEGOTIATE_SECURITY_SIGNATURES_REQUIRE:
            self.log.info('Server requires all SMB messages to be signed')
            self.is_signing_active = (self.sign_options != SMB.SIGN_NEVER)
        elif self.security_mode & NEGOTIATE_SECURITY_SIGNATURES_ENABLE:
            self.log.info('Server supports SMB signing')
            self.is_signing_active = (self.sign_options == SMB.SIGN_WHEN_SUPPORTED)
        else:
            self.is_signing_active = False

        if self.is_signing_active:
            self.log.info('SMB signing activated. All SMB messages will be signed.')
            self.signing_session_key = ntlm_data[-16:]
            if self.capabilities & CAP_EXTENDED_SECURITY:
                self.signing_challenge_response = None
            else:
                self.signing_challenge_response = blob
        else:
            self.log.info('SMB signing deactivated. SMB messages will NOT be signed.')

    def authenticate_SMB2(self, ntlm_data):
        blob = securityblob.generateAuthSecurityBlob(ntlm_data)
        self._sendSMBMessage(SMB2Message(SMB2SessionSetupRequest(blob)))

        if self.security_mode & SMB2_NEGOTIATE_SIGNING_REQUIRED:
            self.log.info('Server requires all SMB messages to be signed')
            self.is_signing_active = (self.sign_options != SMB.SIGN_NEVER)
        elif self.security_mode & SMB2_NEGOTIATE_SIGNING_ENABLED:
            self.log.info('Server supports SMB signing')
            self.is_signing_active = (self.sign_options == SMB.SIGN_WHEN_SUPPORTED)
        else:
            self.is_signing_active = False

        if self.is_signing_active:
            self.log.info('SMB signing activated. All SMB messages will be signed.')
            self.signing_session_key = (ntlm_data[-16:] + b'\0'*16)[:16]
            if self.capabilities & CAP_EXTENDED_SECURITY:
                self.signing_challenge_response = None
            else:
                self.signing_challenge_response = blob
        else:
            self.log.info('SMB signing deactivated. SMB messages will NOT be signed.')

    def _handleNegotiateResponse_SMB1(self, message):
        if message.uid and not self.uid:
            self.uid = message.uid

        if message.hasExtendedSecurity or message.payload.supportsExtendedSecurity:
            ntlm_data = self.ntlm_data
            blob = securityblob.generateNegotiateSecurityBlob(ntlm_data)
            self._sendSMBMessage(SMBMessage(ComSessionSetupAndxRequest__WithSecurityExtension(message.payload.session_key, blob)))

    def _handleNegotiateResponse_SMB2(self, message):
        ntlm_data = self.ntlm_data
        blob = securityblob.generateNegotiateSecurityBlob(ntlm_data)
        self._sendSMBMessage(SMB2Message(SMB2SessionSetupRequest(blob)))

    def _handleSessionChallenge_SMB1(self, message, ntlm_token):
        pass

    def _handleSessionChallenge_SMB2(self, message, ntlm_token):
        pass

class CacheConnections:

    def __init__(self):
        self._mutex = Lock()
        self._cache = {}

    def __len__(self):
        return len(self._cache)

    def remove(self, id):
        self._mutex.acquire()
        (proxy, ts) = self._cache.get(id, (None,None))
        if proxy:
            proxy.close()
            del self._cache[id]
        self._mutex.release()

    def add(self, id, proxy):
        self._mutex.acquire()
        self._cache[id] = ( proxy, int(time.time()) )
        self._mutex.release()

    def cleanTimeout(self):
        now = int(time.time())
        self._mutex.acquire()
        for id, conn in list(self._cache.items()):
            if conn[1]+60<now:
                conn[0].close()
                del self._cache[id]
        self._mutex.release()

    def cleanAll(self):
        now = int(time.time())
        self._mutex.acquire()
        for id, conn in list(self._cache.items()):
            conn[0].close()
            del self._cache[id]
        self._mutex.release()

    def has_key(self,id):
        return id in self._cache

    def get_proxy(self, id):
        self._mutex.acquire()
        proxy = self._cache[id][0]
        self._mutex.release()
        return proxy

cache = CacheConnections()

def decode_http_authorization_header(auth):
    ah = auth.split(' ')
    if len(ah) == 2:
        b64 = base64.b64decode(ah[1])
        if ah[0] == 'NTLM':
            return ('NTLM', b64)
    return False

def dbg(m):
    syslog.syslog(syslog.LOG_ERR, "TEST:" +m)

def ntlm_message_type(msg):
    if not msg.startswith(b'NTLMSSP\x00') or len(msg)<12:
        raise RuntimeError("Not a valid NTLM message: '%s'" % hexlify(msg))
    msg_type = unpack('<I', msg[8:8+4])[0]
    if msg_type not in (1,2,3):
        raise RuntimeError("Incorrect NTLM message Type: %d" % msg_type)
    return msg_type

def parse_ntlm_authenticate(msg):
    NTLMSSP_NEGOTIATE_UNICODE = 0x00000001
    idx = 28
    length, offset = unpack('<HxxI', msg[idx:idx+8])
    domain = msg[offset:offset+length]
    idx += 8
    length, offset = unpack('<HxxI', msg[idx:idx+8])
    username = msg[offset:offset+length]
    idx += 24
    flags = unpack('<I', msg[idx:idx+4])[0]
    if flags & NTLMSSP_NEGOTIATE_UNICODE:
        domain = str(domain.decode('utf-16-le'))
        username = str(username.decode('utf-16-le'))
    return username, domain

def get_connected_dc():
    p = subprocess.Popen(["/usr/local/bin/wbinfo", "-P"], stdout=subprocess.PIPE)
    output, err = p.communicate()
    p_status = p.wait()
    if err:
        raise RuntimeError("wbinfo error: "+err)
    arr = output.split(b"\"")
    if 3 != len(arr):
        raise RuntimeError("wbinfo format wrong: " + output)
    return arr[1]

def process(cid, auth):
    try:
        ah_data = decode_http_authorization_header(auth)
    except:
        ah_data = False

    if not ah_data:
        return b'-1', b"No NTLM authenticate header"

    try:
        ntlm_version = ntlm_message_type(ah_data[1])
        if ntlm_version != 1 and ntlm_version != 3:
            raise RuntimeError('TYPE ' + str(ntlm_version) + ' message in client request.(183)')

        if ntlm_version == 1:
            cache.cleanTimeout()
            proxy = SYNOIWAConnection()
            proxy.connect(get_connected_dc())
            ntlm_token = proxy.negotiate(ah_data[1])
            cache.add(cid, proxy)
            return cid, b"NTLM " + base64.b64encode(bytes(ntlm_token))

        if ntlm_version == 3:
            if not cache.has_key(cid):
                raise RuntimeError("TYPE 3 error: Proxy timeout.(id="+str(cid)+") >60s")
            user, domain = parse_ntlm_authenticate(ah_data[1])
            if not domain:
                raise RuntimeError("user "+ domain + "\\" + user + " invalid.")

            proxy = cache.get_proxy(cid)
            if not proxy.authenticate(ah_data[1]):
                raise RuntimeError('TYPE 3 error: authenticate failed.')
            cache.remove(cid)
            return b'0', domain+"\\"+user

    except Exception as e:
        if cache.has_key(cid): cache.remove(cid)
        return b'-1', str(e)

def run():
    socket_file = "/var/run/synocgid/synopyntlm.htua.ntlm.socket"

    try:
        os.mkdir("/var/run/synocgid", 0o755);
    except OSError:
        pass
    try:
        os.remove(socket_file)
    except OSError:
        pass

    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.bind(socket_file)
    sock.listen(5)
    while True:
        conn, addr = sock.accept()
        recv_data = conn.recv(4096)
        str_data = recv_data.decode()
        if not str_data: continue
        arr = str_data.split(',')
        result = -1
        message = "Format error:" + str_data
        if 0 != os.system("/usr/local/bin/wbinfo -p &> /dev/null"):
            message = "Ping to winbindd failed"
        elif 2 == len(arr):
            result, message = process(arr[0], arr[1])
        conn.send(b(result) + b"," + b(message))
        conn.close()
    try:
        os.remove(socket_file)
    except OSError:
        pass

    cache.cleanAll()

def b(x):
    '''
    porting from PyAuthenNTLM2
    '''
    if str == type(x):
        return codecs.latin_1_encode(x)[0]
    return x

def debug(cid, auth):
    result, message = process(cid, auth)
    print((str(result) + "," + message))

def synontlm(args):
    if args.auth:
        debug(5566, args.auth)
    else:
        run()

parser = argparse.ArgumentParser(description='Process NTLM negotiation.')
parser.add_argument('-a', '--auth', default=False)
args = parser.parse_args()
synontlm(args)