import socket
import sys
import re
import time
import datetime
import argparse
import functools
import traceback


class Simple:
    def __init__(self, host, port, timeout, response_delay, outfile):
        self.__dict__.update(vars())
        del self.self

    @classmethod
    def start(cls):
        options = cls.get_args()
        server = cls(**vars(options))
        server.serve()

    @staticmethod
    def get_args():
        p = argparse.ArgumentParser(conflict_handler="resolve")
        p.add_argument('-h', '--host', help="Bind to IP address", default='')
        p.add_argument('-p', '--port', type=int, help="Bind to port", default=80)
        p.add_argument('-t', '--timeout', type=int, help="Socket timeout", default=3)
        p.add_argument(
            '-o',
            '--outfile',
            default=sys.stdout,
            type=functools.partial(open, mode='wb'),
            help='save output to file',
        )

        def seconds(seconds):
            return datetime.timedelta(seconds=seconds)

        p.add_argument(
            '-d',
            '--delay',
            dest='response_delay',
            type=seconds,
            help="Artificial delay in response",
            default=datetime.timedelta(),
        )
        return p.parse_args()

    def serve_one(self):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.bind((self.host, self.port))
        s.listen(1)
        self.conn, addr = s.accept()
        print('Accepted connection from', addr)
        self.get_response()

    serve = serve_one

    def get_response(self):
        try:
            self.conn.settimeout(self.timeout)
            headers, content = self.get_headers(self.conn)
            content_len = self.get_content_length(headers) or 0
            content = self.get_content(self.conn, content, content_len)
            self.conn.send(b'HTTP/1.0 200 OK\r\n')
            time.sleep(self.response_delay.total_seconds())
            self.conn.send(b'\r\nGot It!')
        except socket.error as e:
            print('Error %s' % e)
            if content:
                print('partial result')
                print(repr(content), file=self.outfile)
        finally:
            self.conn.close()

    @staticmethod
    def get_content_length(request):
        match = re.search(r'^Content-Length:\s+(\d+)\s*$', request, re.I | re.MULTILINE)
        if match:
            return int(match.group(1))
        print('no content length found', file=sys.stderr)

    @staticmethod
    def get_headers(conn):
        res = b''
        while b'\r\n\r\n' not in res:
            res += conn.recv(1024)
        bytes = len(res)
        headers, _sep, content = res.partition(b'\r\n\r\n')
        print('received', bytes, 'bytes', file=sys.stderr)
        headers = headers.decode('utf-8')
        print(headers)
        return headers, content

    def get_content(self, conn, content='', length=0):
        while len(content) < length:
            content += conn.recv(1024)
        bytes = len(content)
        print('received', bytes, 'bytes content', file=sys.stderr)
        print(content, file=self.outfile)
        return content


class AuthRequest(Simple):
    def serve_until_auth(self):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.bind((self.host, self.port))
        s.listen(1)
        while True:
            self.conn, addr = s.accept()
            print(f'Accepted connection from {addr}')

            if not self.check_auth_response(self.conn) == 'retry':
                break

    serve = serve_until_auth

    def check_auth_response(self, conn):
        try:
            conn.settimeout(self.timeout)
            headers, content = self.get_headers(conn)
            content_len = self.get_content_length(headers) or 0
            content = self.get_content(self.conn, content, content_len)
            user_pat = re.compile(r'^Authorization:\s+(.*)\s*$', re.I | re.MULTILINE)
            matched_header = user_pat.search(headers)
            if matched_header:
                conn.send(b'HTTP/1.0 200 OK\r\n')
                user = matched_header.group(1).encode('ascii')
                conn.send(b'\r\nYou are authenticated as ' + user)
                return
            conn.send(b'HTTP/1.0 401 Authorization Required\r\n')
            conn.send(b'Connection: close\r\n')
            msg = b'Go get me some credentials'
            conn.send(b'Content-Length: %d\r\n' % len(msg))
            conn.send(b'WWW-Authenticate: NTLM\r\n')
            conn.send(b'WWW-Authenticate: Basic realm="fake-auth"\r\n')
            conn.send(b'\r\n')
            conn.send(msg)
            print('sent authorization request')
        except socket.error:
            print('error in connection')
            traceback.print_exc()
        finally:
            conn.close()
        return 'retry'
