from __future__ import print_function
from opencanary.modules import CanaryService
from opencanary.config import ConfigException

from twisted.protocols.policies import TimeoutMixin
from twisted.internet.protocol import Protocol
from twisted.internet.protocol import Factory
from twisted.application import internet
from ntlmlib.messages import ChallengeResponse, TargetInfo

import struct
import re
import collections


# Monkeypatch bug in ntmllib
if getattr(TargetInfo, 'getData', None) is None:
    def getData(self):
        return self.get_data()
    TargetInfo.getData = getData

TDSPacket = collections.namedtuple('TDSPacket', 'type status spid packetid window payload')
PreLoginOption = collections.namedtuple('PreLoginOption', 'token data')

class MSSQLProtocol(Protocol, TimeoutMixin):
    # overview https://msdn.microsoft.com/en-us/library/dd357422.aspx

    TDS_HEADER_LEN = 8
    TDS_TYPE_PRELOGIN = 0x12
    TDS_TYPE_RESPONSE = 0x04
    TDS_TYPE_LOGIN7 = 0x10
    TDS_TYPE_SSPI = 0x11

    # https://msdn.microsoft.com/en-us/library/dd357559.aspx
    PRELOGIN_VERSION = 0x00
    PRELOGIN_ENCRYPTION = 0x01
    PRELOGIN_INSTOPT = 0x02
    PRELOGIN_THREADID = 0x03
    PRELOGIN_MARS = 0x04
    PRELOGIN_TRACEID = 0x05
    PRELOGIN_FEDAUTHREQUIRED = 0x06
    PRELOGIN_NONCEOPT = 0x07

    # https://msdn.microsoft.com/en-us/library/dd304019.aspx
    LOGIN7_OPT1_BYTEORDER_MASK = 0x01
    LOGIN7_OPT1_BYTEORDER_X86 = 0x01
    LOGIN7_OPT1_BYTEORDER_68000 = 0x01
    LOGIN7_FIELDS = ["length", "TDSversion", "packetSize", "clientProgVer", "clientPID", "connID", "optFlags1", "optFlags2", "typeFlags", "optFlags3", "clientTimeZone", "clientLCID", "ibHostName", "cchHostName", "ibUserName", "cchUserName", "ibPassword", "cchPassword", "ibAppName", "cchAppName", "ibServerName", "cchServerName", "ibExtension", "cbExtension", "ibCltIntName", "cchCltIntName", "ibLanguage", "cchLanguage", "ibDatabase", "cchDatabase", "ClientID", "ibSSPI", "cbSSPI", "ibAtchDBFile", "cchAtchDBFile", "ibChangePassword", "cchChangePassword", "cbSSPILong"]

    NMAP_PROBE_1 = TDSPacket(
        type=18,
        status=1,
        spid=0,
        packetid=0,
        window=0,
        payload=b'\x00\x00\x15\x00\x06\x01\x00\x1b\x00\x01\x02\x00\x1c\x00\x0c\x03\x00(\x00\x04\xff\x08\x00\x01U\x00\x00\x00MSSQLServer\x00H\x0f\x00\x00')

    NMAP_PROBE_1_RESP = {
        "2008R2" : b"\x04\x01\x00.\x00\x00\x01\x00\x00\x00\x15\x00\x06\x01\x00\x1b\x00\x01\x02\x00\x1c\x00\x01\x03\x00\x1d\x00\x00\xff\x0a\x32\x10\xb4",
        "2012" : b"\x04\x01\x00\x25\x00\x00\x01\x00\x00\x00\x15\x00\x06\x01\x00\x1b\x00\x01\x02\x00\x1c\x00\x01\x03\x00\x1d\x00\x00\xff\x0b\x00\x0c\x38",
        "2014" : b"\x04\x01\x00\x25\x00\x00\x01\x00\x00\x00\x15\x00\x06\x01\x00\x1b\x00\x01\x02\x00\x1c\x00\x01\x03\x00\x1d\x00\x00\xff\x0c\x00\x07\xd0"
    }

    def __init__(self, factory):
        self._busyReceiving = False
        self._buffer = b""
        self.factory = factory
        self.setTimeout(10)

    @staticmethod
    def build_packet(tds):
        header = struct.pack('>BBHHBB',
                             tds.type,
                             tds.status,
                             len(tds.payload) + MSSQLProtocol.TDS_HEADER_LEN,
                             tds.spid,
                             tds.packetid,
                             tds.window
        )

        return  header + tds.payload

    @staticmethod
    def parsePreLogin(data):
        # parse the initial packet
        index = data.find("\xff")

        if index <= 0:
            return None

        options = data[:index]

        if (len(options) % 5 != 0):
            return None

        def getOption(i):
            (token, offset, length) = struct.unpack(">BHH", options[i:i+5])
            tokendata = data[offset: offset + length]
            assert(len(tokendata) == length)
            return PreLoginOption._make((token, tokendata))

        try:
            return map(getOption, range(0, len(options), 5))
        except Exception as e:
            print(e)

        return None

    @staticmethod
    def buildPreLogin(preloginopts):
        preloginopts.sort(key=lambda x: x.token)
        dataoffset = len(preloginopts) * 5 + len("\xff")
        data = b""
        header = b""

        for opt in preloginopts:
            token = opt.token
            offset = dataoffset + len(data)
            length = len(opt.data)
            header += struct.pack(">BHH", token, offset, length)
            data += opt.data

        return header + b"\xff" +  data

    @staticmethod
    def parseLogin7(data):
        # These header values apper to be little endian, even though
        # the byte-order flag suggests otherwise. In case it is
        # dependent, simply check the endian flag in optFlags1 in
        # before parsing the rest in ord(data[24]). For the moment,
        # we'll assume that flag applies only to the data part.
        hfmt = "< 6I 4B l I 18H 6s 6H I"
        hlen = struct.calcsize(hfmt)
        try:
            htuple = struct.unpack(hfmt, data[:hlen])
        except Exception as e:
            return None

        def decodePassChar(c):
            # https://msdn.microsoft.com/en-us/library/dd304019.aspx
            c = ord(c) ^ 0xa5