<reponame>Tombmyst/Empire
from empire.python.typings import *
import requests
from os.path import basename

from empire.http_enum.http_verbs import HTTPVerbs
from empire.http_enum.http_headers import HTTPHeaders

from empire.request_util.tomb_response import TombResponse
from empire.util.log import *
from empire.util.execution_report import the_report


class RequestUtil:
    MY_TOKEN: Final[str] = '<KEY>'
    HTTP_VERB_MAPPING: Dict[str, Callable] = {
        HTTPVerbs.POST: requests.post,
        HTTPVerbs.GET: requests.get,
        HTTPVerbs.PUT: requests.put,
        HTTPVerbs.PATCH: requests.patch,
        HTTPVerbs.HEAD: requests.head,
        HTTPVerbs.DELETE: requests.delete,
        HTTPVerbs.OPTIONS: requests.options
    }

    @staticmethod
    def send_file_to_presigned_s3_url(url: str, fields: JSON, file_name: str) -> Union[TombResponse, None]:
        with open(file_name, 'rb') as f:
            files = {'file': (basename(file_name), f)}
            post_response = requests.post(
                url,
                data=fields,
                files=files
            )
            return TombResponse(post_response, take_body_from_text=True)
            #print(post_response.status_code, post_response.text)
        # with open(file_name, 'rb') as f:
        #     files = {'file': (basename(file_name), f)}
        #     return RequestUtil.call_api((HTTPVerbs.POST, url), data=fields, files=files, take_body_from_text=True)

    @staticmethod
    def call_api(endpoint_data: Tuple[str, str],
                 body: Union[JSON, None] = None,
                 data: Union[JSON, None] = None,
                 files: Union[Dict[str, Tuple[str, BinaryIO]], None] = None,
                 additional_headers: Union[Dict[str, str], None] = None,
                 query_parameters: Union[Dict[str, str], None] = None,
                 send_binary_response_to: Union[str, None] = None,
                 ignore_body: bool = False,
                 take_body_from_text: bool = False) -> Union[TombResponse, None]:
        if not endpoint_data or len(endpoint_data) != 2 or endpoint_data[0] not in RequestUtil.HTTP_VERB_MAPPING:
            return None

        params: Dict[str, Any] = {
            'headers': {
                HTTPHeaders.AUTHORIZATION: f'BEARER {RequestUtil.MY_TOKEN}',
                HTTPHeaders.CONTENT_TYPE: 'application/json'
            }
        }

        if body is not None:
            params['json'] = body

        if data is not None:
            params['data'] = data

        if files is not None:
            params['files'] = files

        if additional_headers:
            params['headers'].update(additional_headers)

        the_url: str = endpoint_data[1]
        if query_parameters is not None:
            the_url += '?'
            for query_parameter_name in query_parameters.keys():
                the_url += '{}={}&'.format(query_parameter_name, query_parameters[query_parameter_name])
            the_url = the_url[:-1]

        response: requests.Response = RequestUtil.HTTP_VERB_MAPPING[endpoint_data[0]](url=the_url, **params)

        if send_binary_response_to:
            with open(send_binary_response_to, 'wb') as f:
                print('Writing data to:', send_binary_response_to)
                f.write(response.content)
            the_report.add_step('Doing API call {}'.format(str(endpoint_data)))
            return TombResponse(response, ignore_body=True)

        the_report.add_step('Doing API call {}'.format(str(endpoint_data)))
        return TombResponse(response, ignore_body=ignore_body, take_body_from_text=take_body_from_text)

    @staticmethod
    def download_by_streaming(url: str, download_to: str, is_binary: bool = True) -> bool:
        try:
            Log.info('Starting download from {}'.format(url), __file__, get_function_name())
            with requests.get(url, stream=True) as r:
                r.raise_for_status()
                with open(download_to, 'wb' if is_binary else 'w') as f:
                    for chunk in r.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
            return True
        except Exception as error:
            Log.error('An error occurred while streaming data', __file__, get_function_name(), error=error)
            return False

from empire.python.typings import *
import requests
from os.path import basename

from empire.http_enum.http_verbs import HTTPVerbs
from empire.http_enum.http_headers import HTTPHeaders

from empire.request_util.tomb_response import TombResponse
from empire.util.log import *
from empire.util.execution_report import the_report


class RequestUtil:
    MY_TOKEN: Final[str] = '<KEY>'
    HTTP_VERB_MAPPING: Dict[str, Callable] = {
        HTTPVerbs.POST: requests.post,
        HTTPVerbs.GET: requests.get,
        HTTPVerbs.PUT: requests.put,
        HTTPVerbs.PATCH: requests.patch,
        HTTPVerbs.HEAD: requests.head,
        HTTPVerbs.DELETE: requests.delete,
        HTTPVerbs.OPTIONS: requests.options
    }

    @staticmethod
    def send_file_to_presigned_s3_url(url: str, fields: JSON, file_name: str) -> Union[TombResponse, None]:
        with open(file_name, 'rb') as f:
            files = {'file': (basename(file_name), f)}
            post_response = requests.post(
                url,
                data=fields,
                files=files
            )
            return TombResponse(post_response, take_body_from_text=True)
            #print(post_response.status_code, post_response.text)
        # with open(file_name, 'rb') as f:
        #     files = {'file': (basename(file_name), f)}
        #     return RequestUtil.call_api((HTTPVerbs.POST, url), data=fields, files=files, take_body_from_text=True)

    @staticmethod
    def call_api(endpoint_data: Tuple[str, str],
                 body: Union[JSON, None] = None,
                 data: Union[JSON, None] = None,
                 files: Union[Dict[str, Tuple[str, BinaryIO]], None] = None,
                 additional_headers: Union[Dict[str, str], None] = None,
                 query_parameters: Union[Dict[str, str], None] = None,
                 send_binary_response_to: Union[str, None] = None,
                 ignore_body: bool = False,
                 take_body_from_text: bool = False) -> Union[TombResponse, None]:
        if not endpoint_data or len(endpoint_data) != 2 or endpoint_data[0] not in RequestUtil.HTTP_VERB_MAPPING:
            return None

        params: Dict[str, Any] = {
            'headers': {
                HTTPHeaders.AUTHORIZATION: f'BEARER {RequestUtil.MY_TOKEN}',
                HTTPHeaders.CONTENT_TYPE: 'application/json'
            }
        }

        if body is not None:
            params['json'] = body

        if data is not None:
            params['data'] = data

        if files is not None:
            params['files'] = files

        if additional_headers:
            params['headers'].update(additional_headers)

        the_url: str = endpoint_data[1]
        if query_parameters is not None:
            the_url += '?'
            for query_parameter_name in query_parameters.keys():
                the_url += '{}={}&'.format(query_parameter_name, query_parameters[query_parameter_name])
            the_url = the_url[:-1]

        response: requests.Response = RequestUtil.HTTP_VERB_MAPPING[endpoint_data[0]](url=the_url, **params)

        if send_binary_response_to:
            with open(send_binary_response_to, 'wb') as f:
                print('Writing data to:', send_binary_response_to)
                f.write(response.content)
            the_report.add_step('Doing API call {}'.format(str(endpoint_data)))
            return TombResponse(response, ignore_body=True)

        the_report.add_step('Doing API call {}'.format(str(endpoint_data)))
        return TombResponse(response, ignore_body=ignore_body, take_body_from_text=take_body_from_text)

    @staticmethod
    def download_by_streaming(url: str, download_to: str, is_binary: bool = True) -> bool:
        try:
            Log.info('Starting download from {}'.format(url), __file__, get_function_name())
            with requests.get(url, stream=True) as r:
                r.raise_for_status()
                with open(download_to, 'wb' if is_binary else 'w') as f:
                    for chunk in r.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
            return True
        except Exception as error:
            Log.error('An error occurred while streaming data', __file__, get_function_name(), error=error)
            return False
#print(post_response.status_code, post_response.text) # with open(file_name, 'rb') as f: #     files = {'file': (basename(file_name), f)} #     return RequestUtil.call_api((HTTPVerbs.POST, url), data=fields, files=files, take_body_from_text=True)