#!/usr/bin/env python
# encoding: utf8

"""
universal http request representation format

same as har request representation

spec: http://www.softwareishard.com/blog/har-12-spec/

method: GET, POST

URL:

headers:
[
    {
        "name": "Accept-Encoding",
        "value": "gzip,deflate",
        "comment": ""
    },
    {
        "name": "Accept-Language",
        "value": "en-us,en;q=0.5",
        "comment": ""
    }
]

cookies:
[
    {
        "name": "TestCookie",
        "value": "Cookie Value",
        "path": "/",
        "domain": "www.janodvarko.cz",
        "expires": "2009-07-24T19:20:30.123+02:00",
        "httpOnly": false,
        "secure": false,
        "comment": ""
    }
]

postData is not in specification, here we use bytes
    so be it, if met any other form, transform it.


be careful when reuse some headers, remember to remove Content-Length.
"""

import os
import cgi
import json
import random
import requests
import copy
import urllib
import urlparse
import traceback
import tldextract
try:
    from http_parser.parser import HttpParser
except ImportError:
    from http_parser.pyparser import HttpParser

from pprint import pprint
from j_pentest_pyobj import objwalk, objset

class Request(object):

    def __init__(self, d):
        """
            d: dict
            keys required: method, url, headers, cookies 
            keys optional: postData, response
        """
        self.d = d # store raw data
        self.method = d['method']
        self.url = d['url']
        self.headers = d.get( 'headers', [] )
        self.cookies = d.get( 'cookies', [] )
        self.postData = d.get( 'postData', '' )

        # my fields
        self.tld = tldextract.extract( self.url ).registered_domain

        self.url_parts = urlparse.urlparse( self.url )
        self.path_components = self.url_parts.path.split( '/' )
        self.query_components = urlparse.parse_qsl( self.url_parts.query )
        # notice: i used parse_qsl instead of parse_qs ( notice the l )
        # example of query_components:  [('wd', 'wtf'), ('wd2', ''), ('wd', 'wtf2')]

        # content-length is not concerned, and actually it would affect fuzz, so remove it
        self.headers = filter( lambda header:header['name'].lower()!='content-length', self.headers )
        # cookies are in cookies field, so remove it from headers
        self.headers = filter( lambda header:header['name'].lower()!='cookie', self.headers )

    def __iter__( self ):
        yield 'method', self.method
        yield 'url', self.url
        yield 'headers', self.headers
        yield 'cookies', self.cookies
        yield 'postData', self.postData

    def _key( self ): # hash key
        if self.method.lower() == 'get' or self.method.lower() == 'post':
            u = urlparse.urlparse( self.url )
            q = tuple( urlparse.parse_qs( u.query ).keys() )
            return ( u.netloc, u.path, q )
        else:
            return ( self.method, self.url )
    
    def __hash__( self ):
        return hash( self._key() )

    def request( self ):
        headers = dict( map( lambda header: ( header['name'], header['value'] ), self.headers ) )
        cookies = dict( map( lambda cookie: ( cookie['name'], cookie['value'] ), self.cookies ) )
        if self.method.lower() == 'get':
            data = ''
        else:
            data = self.postData
        r  =  requests.request( method = self.method, 
                                url = self.url,
                                headers = headers,
                                cookies = cookies,
                                data = data,
                                allow_redirects = True,
                                timeout = 7,
                              )
        return r

    def from_raw( self ):
        # todo, use HttpParser to parse raw request
        pass

    def to_raw( self ):
        template = u'{method} {url} HTTP/1.1\r\n{headers}\r\n\r\n{postData}'

        cookies = copy.deepcopy( self.cookies )
        cookies = '; '.join( map( lambda cookie:'{name}={value}'.format(name=cookie['name'], value=cookie['value'], ), cookies ) )

        headers = copy.deepcopy( self.headers )
        headers = '\r\n'.join( map( lambda header:'{name}: {value}'.format(name=header['name'], value=header['value'], ), headers ) )
        if cookies:
            headers += '\r\n' + 'Cookie: ' + cookies

        return template.format( method=self.method, url=self.url, headers=headers, postData=sefl.postData )


    def to_tagged_raw( self ):
        tag = u'\u2694' # crossed swords :)
        template = u'{method} {url} HTTP/1.1\r\n{headers}\r\n\r\n{postData}'

        url_parts = list( copy.deepcopy( self.url_parts ) )
        new_path = '/'.join( [ '{tag}{c}{tag}'.format(tag=tag,c=c) if c else '' for c in self.path_components ] )
        url_parts[2] = new_path
        new_query = '&'.join( [ '{key}={tag}{value}{tag}'.format( key=key, tag=tag, value=value ) for key,value in self.query_components ] )
        url_parts[-2] = new_query
        url = urlparse.urlunparse( url_parts )

        cookies = copy.deepcopy( self.cookies )
        cookies = '; '.join( map( lambda cookie:'{name}={tag}{value}{tag}'.format(name=cookie['name'], value=cookie['value'], tag=tag), cookies ) )

        headers = copy.deepcopy( self.headers )
        headers = '\r\n'.join( map( lambda header:'{name}: {tag}{value}{tag}'.format(name=header['name'], value=header['value'], tag=tag), headers ) )
        if cookies:
            headers += '\r\n' + 'Cookie: ' + cookies

        postData = copy.deepcopy( self.postData )

        return template.format( method=self.method, url=url, headers=headers, postData=postData )

    def fuzz( self, payloads, fuzz_url=True, fuzz_query=True, fuzz_headers=True, fuzz_cookies=True, fuzz_postData=True ):
        """
            naive fuzz, fuzz any fuzzable field
            url: urlpath
            query: url query params
            headers: header fields
            cookies: cookies
            postData: form fields, json fields
        """
        # fuzz url path components
        if fuzz_url:
            path_components = copy.deepcopy( self.path_components )

            for path_component_index in xrange(len(path_components)):

                for payload in payloads:
                    
                    if path_components[path_component_index] == '':
                        break

                    new_path_components = copy.deepcopy( self.path_components )
                    new_path_components[path_component_index] = path_components[path_component_index] + payload
                    new_path = '/'.join( new_path_components )

                    new_url_parts = list( copy.deepcopy( self.url_parts ) )
                    new_url_parts[2] = new_path
                    pprint( new_url_parts )
                    new_url = urlparse.urlunparse( new_url_parts )

                    d = copy.deepcopy( self.d )
                    d['url'] = new_url

                    yield Request( d )
                    
        # fuzz url query parameters
        if fuzz_query:
            query_components = copy.deepcopy( self.query_components )

            for query_components_index in xrange(len(query_components)):

                for payload in payloads:

                    new_query_components = copy.deepcopy( self.query_components )
                    new_query_components = map( lambda x:list(x), new_query_components )
                    new_query_components[query_components_index][1] = query_components[query_components_index][1] + payload
                    new_query = '&'.join( '{}={}'.format( urllib.quote(key), urllib.quote(value) ) for key,value in new_query_components )

                    new_url_parts = list( copy.deepcopy( self.url_parts ) )
                    new_url_parts[-2] = new_query
                    new_url = urlparse.urlunparse( new_url_parts )

                    d = copy.deepcopy( self.d )
                    d['url'] = new_url

                    yield Request( d )
                
        # fuzz headers
        if fuzz_headers:
            for header_index in xrange(len(self.headers)):

                for payload in payloads:

                    new_headers = copy.deepcopy( self.headers )
                    new_headers[header_index]['value'] = self.headers[header_index]['value'] + payload

                    d = copy.deepcopy( self.d )
                    d['headers'] = new_headers

                    yield Request( d )
                
        # fuzz cookies
        if fuzz_cookies:
            for cookie_index in xrange(len(self.cookies)):

                for payload in payloads:

                    new_cookies = copy.deepcopy( self.cookies )
                    new_cookies[cookie_index]['value'] = self.cookies[cookie_index]['value'] + payload
                    
                    d = copy.deepcopy( self.d )
                    d['cookies'] = new_cookies

                    yield Request( d )
                
        # fuzz postData
        if fuzz_postData:
            content_type = self.infer_content_type()
            if content_type == 'application/x-www-form-urlencoded':
                postData = copy.deepcopy( self.postData )
                postData = urlparse.parse_qsl( postData )
                pprint( postData )
                for kv_index in xrange(len(postData)):
                    for payload in payloads:
                        new_postData = copy.deepcopy( postData )
                        new_postData = map( lambda x:list(x), new_postData )
                        new_postData[kv_index][1] = new_postData[kv_index][1] + payload
                        new_postData = '&'.join( '{}={}'.format( urllib.quote(k), urllib.quote(v) ) for k,v in new_postData )

                        d = copy.deepcopy( self.d )
                        d['postData'] = new_postData

                        yield Request(d)

            elif content_type == 'application/json':
                for path, val in objwalk( json.loads( self.postData ) ):
                    for payload in payloads:
                        new_postData = json.loads( copy.deepcopy( self.postData ) )

                        # todo: more precise control over field
                        objset( new_postData, path, str(val)+payload )

                        d = copy.deepcopy( self.d )
                        d['postData'] = json.dumps( new_postData )

                        yield Request(d)

            # todo: more post data type fuzz

            else:
                new_postData = copy.deepcopy( self.postData ) + payload

                d = copy.deepcopy( self.d )
                d['postData'] = new_postData

                yield Request(d)

    def infer_content_type( self ):
        supported_content_type = [ 'application/x-www-form-urlencoded', 'application/json' ]
        content_type_header = filter( lambda header:header['name'].lower()=='content-type', self.headers )
        if content_type_header:
            content_type = cgi.parse_header( content_type_header[0]['value'] )[0]
            if content_type in supported_content_type:
                return content_type
        return 'raw'

    @staticmethod
    def from_mitm_flow( flow ):
        """
            generate Request from mitm flow
            mitm flow headers format:
            {
                'accept-encoding': 'gzip',
                ....
            }
            mitm flow cookies format( after converted to dict ):
            {
                'gid': [ '14235665554', ],
            }
        """
        headers = dict( flow['request_headers'] )
        headers = map( lambda key: {'name':key, 'value':headers[key]}, headers.keys() )

        _cookies = dict( flow['cookies'] )
        cookies = []
        for key in _cookies:
            for val in _cookies[key]:
                cookies.append( {'name':key, 'value':val} )

        d = {
            'method': flow['method'],
            'url': flow['url'],
            'cookies': cookies,
            'headers': headers,
            'postData': flow['request_content'],
        }
        return Request( d )

    def to_sqlmap( self ):
        """
            generate sqlmapapi options
        """
        url = urlparse.urlparse( self.url )
        path = url.path.split( '/' )
        # todo: path injection

        headers = '\r\n'.join( map( lambda header:'{}: {}'.format( header['name'], header['value'] ), headers ) )
        options = {
            'level': '5',
            'risk': '1',
            'threads': '3',
            'smart': True,

            'url': self.url,
            'method': self.method,
            'headers': headers,
            'data': self.postData,
        }

        return options
    
    def is_media( self ):
        """
            from path extract extension and guess file type based on that
        """
        media_ext = "flv,mp4,mp4,swf,jpg,jpeg,png,mp4,gif,pdf,rar,zip,avi,mp4,swf,wmi,exe,mpeg,ppt,pptx,doc,docx,xls,xlsx,apk"
        media_ext = filter( None, media_ext.split(',') )
        ext = self.get_ext()
        return ( (ext) and ( ext in media_ext ) )

    def should_sqli( self ):
        """
            from extension to determin whether should be test against sqli, static data no need to test
        """
        media_ext = "js,flv,mp4,mp4,swf,jpg,jpeg,png,mp4,gif,pdf,rar,zip,avi,mp4,swf,wmi,exe,mpeg,ppt,pptx,doc,docx,xls,xlsx,apk"
        media_ext = filter( None, media_ext.split(',') )
        ext = self.get_ext()
        return ( ext not in media_ext )

    def get_ext( self ):
        """
            get extension from url, so this is not reliable
            os.path.splitext is used to extract extension
        """
        path = urlparse.urlparse( self.url ).path
        ext = os.path.splitext(path)[1]
        return ext




