# _*_ coding:utf-8 _*_

import re
from urllib import parse
from urllib.parse import urlencode

from pentest.net.socket import SocketClient


class HttpClient(object):
    

    def __init__(self):

        # Local memory storage for cookies
        self.cookies  = {}

        # Last url for sent the referer header
        self.last_url = None

        # Socket helper handler
        self.socket_client = SocketClient()

        # Handle events
        self.socket_client.setOnReceive(self.onReceive)
        self.socket_client.setOnConnect(self.onConnect)

        # Current request context
        self.current_context = {
            'hostname'       : None,
            'port'           : None,
            'use-ssl'        : None,
            'url'            : None,
            'uri'            : None,
            'custom-headers' : None,
            'post-data'      : None,
            'raw-request'    : None,
            'raw-response'   : None
        }


    def clearContext(self):

        self.cookies  = {}
        self.last_url = None


    def request(self, url, custom_headers=None, post_data=None):
        
        # to bytes array
        if(not isinstance(url, bytes)):
            url = str(url).encode('utf-8', 'ignore')

        if(post_data):
            
            if(isinstance(post_data, dict)):
                post_data = urlencode(post_data).encode('utf-8', 'ignore')

            if(isinstance(post_data, str)):
                post_data = post_data.encode('utf-8', 'ignore')

        # if(custom_headers):
        #     # Join custom headers
        #     headers.update(custom_headers)

        # Parse URL
        url_parsed = parse.urlparse(url)
        url_data   = {
            'original' : url,
            'path'     : url_parsed.path,
            'host'     : url_parsed.netloc,
            'port'     : url_parsed.port,
            'scheme'   : url_parsed.scheme,
            'query'    : url_parsed.query,
            'uri'      : url_parsed.path + ((b'?' + url_parsed.query) if url_parsed.query else b'')
        }

        if url_data['path'] == b'':
            url_data['path'] = b'/'

        if(not url_data['port']):
            url_data['port'] = 443 if (url_data['scheme'] == b'https') else 80
        
        # Strip the custom port from the hostname
        if(b':' in url_data['host']):
            url_data['host'] = url_data['host'].split(b':')[0]

        # Copy arguments to current context
        self.current_context['hostname']       = url_data['host'].decode()
        self.current_context['port']           = int(url_data['port'])
        self.current_context['use-ssl']        = bool(url_data['scheme'] == b'https')
        self.current_context['url']            = str(url)
        self.current_context['uri']            = url_data['uri'].decode()
        self.current_context['custom-headers'] = custom_headers
        self.current_context['post-data']      = post_data
        self.current_context['raw-response']   = b''

        # Connect to server
        self.socket_client.connect(
            hostname = self.current_context['hostname'],
            port     = self.current_context['port'],
            use_ssl  = self.current_context['use-ssl']
        ) # Locked for connection

        # The socket is closed in this point

        # Save the cookies for next request
        matches = re.findall(
            br'set\-cookie:\s*(.*?);',
            self.current_context['raw-response'],
            re.IGNORECASE | re.MULTILINE
        )

        if len(matches) > 0:
            for cookie in matches:

                cookie = re.search(
                    rb'(.*?)=(.*)',
                    cookie,
                    re.IGNORECASE | re.MULTILINE
                )
                
                var = cookie.group(1).strip()
                val = cookie.group(2).strip()

                if val:
                    self.cookies[var] = val

                else:
                    # Delete cookie if exist
                    if var in self.cookies.keys():
                        self.cookies.pop(var, None)

        # Get the HTTP status code
        status_code = 0
        matches     = re.search(
            br'HTTP\/\d\.\d (\d+) ',
            self.current_context['raw-response'],
            re.IGNORECASE | re.MULTILINE
        )
        if(matches):
            status_code = int(matches.group(1))

        # Split the header and body
        body    = self.current_context['raw-response'].split(b'\r\n\r\n')
        headers = body.pop(0).strip()
        body    = b'\r\n'.join(body)

        # Free memory
        self.current_context['raw-response'] = b''

        # Split headers
        if(b'\r\n' in headers):
            tmp = {}
            for item in headers.split(b'\r\n'):
                value    = item.split(b':')
                key      = value.pop(0).strip()
                value    = b':'.join(value).strip()
                tmp[key] = value
            headers = tmp

        # Parse the raw body

        # TODO: Add better support
        # Transfer-Encoding: chunked  [OK]
        # Transfer-Encoding: compress
        # Transfer-Encoding: deflate
        # Transfer-Encoding: gzip
        # Transfer-Encoding: identity

        if(b'Transfer-Encoding' in headers):
        
            # Chunked response body?
            # https://tools.ietf.org/rfc/rfc7230.txt
            if(headers[b'Transfer-Encoding'] == b'chunked'):

                # Length in hex, convert to int base 16
                bytes_length = b''
                buffer_body  = b''
                body         = b'\r\n' + body # Padding left for bytes length

                while(True):

                    # [\r\n] [length] [\r\n] [part]

                    groups       = re.search(br'\r\n(\w+?)\r\n', body)
                    bytes_length = groups.group(1)
                    bytes_length = int(bytes_length, 16) # Convert length value
                    body         = body[len(groups.group(0)):]# Strip the length string

                    if(bytes_length == 0):
                        break
                    
                    buffer_body      += body[:bytes_length]
                    body             =  body[bytes_length:]

                # Transfer value
                body = buffer_body

                # Free memory
                buffer_body = None

        # Return the result
        return {
            'status-code'      : status_code,
            'response-content' : body,
            'response-headers' : headers,
            'request-content'  : self.current_context['raw-request']
        }


    def onConnect(self, hostname, port, use_ssl):
        
        # Get the cookies
        cookies = []

        for key, value in self.cookies.items():
            cookies.append(key + b'=' + value)

        cookies = b'; '.join(cookies)

        # Begin of packet to send
        self.current_context['raw-request'] = b'\r\n'.join([
            (b'POST ' if self.current_context['post-data'] else b'GET ') + self.current_context['uri'].encode() + b' HTTP/1.1',
            b'Host: ' + self.current_context['hostname'].encode(),
            b'User-Agent: Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:65.0) Gecko/20100101 Firefox/54.0',
            b'Accept: text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
            b'Accept-Language: en-US',
            b'Referer: ' + (self.last_url.encode() if self.last_url else self.current_context['url'].encode()),
            b'Cookie: ' + cookies,
            b'Connection: close'
        ]) + b'\r\n'

        # Parse custom headers
        if(self.current_context['custom-headers']):
            custom_headers = []

            for key, value in self.current_context['custom-headers'].items():
                custom_headers.append(key.encode() + b': ' + value.encode('utf-8', errors='ignored'))

            custom_headers = b'\r\n'.join(custom_headers)
            self.current_context['raw-request'] += custom_headers + b'\r\n'

        # Send post data?
        if(self.current_context['post-data']):
            self.current_context['raw-request'] += b'\r\n'.join([
                b'Content-Type: application/x-www-form-urlencoded',
                b'Content-Length: ' + str(len(self.current_context['post-data'])).encode(),
                b'',
                self.current_context['post-data']
            ])
        else:
            self.current_context['raw-request'] += b'\r\n'

        # Save the url for sent into next request
        self.last_url = self.current_context['url']

        # Send binary packet to server
        self.socket_client.send(self.current_context['raw-request'])


    def onReceive(self, binary_data):
        
        # Append to buffer
        self.current_context['raw-response'] += binary_data
