<commit_before>import logging
import re

import requests
import httpretty

NOT_FOUND = '404 - Not Found'


def mangle(method, uri, headers):
    '''
    See help at common_arguments.add_mangle_arguments to understand what this is.
    '''
    # 'http://ec2-54-254-24-239.ap-southeast-1.compute.amazonaws.com/?url=%s'
    VULN_URL = None
    
    mangled_url = VULN_URL % uri
    
    logging.debug('Requesting %s' % mangled_url)
    response = requests.get(mangled_url)
    
    code = 200
    if NOT_FOUND in response.text:
        code = 404
        
    return (code, headers, response.text)
    
def setup_mangle(cmd_args):
    if cmd_args.mangle_function is not None:
        # We need to verify that the user configured function exists
        try:
            mangle_module_funct = cmd_args.mangle_function
            mangle_module = '.'.join(mangle_module_funct.split('.')[:-1])
            mangle_funct = mangle_module_funct.split('.')[-1]
            
            _temp = __import__(mangle_module, globals(), locals(), ['mangle_funct'], -1)
            mangle_funct = getattr(_temp, mangle_funct)
        except:
            msg = 'The user configured mangle function %s does not exist.'
            logging.critical(msg % mangle_funct)
        
        # enable HTTPretty so that it will monkey patch the socket module
        httpretty.enable()
    
        httpretty.register_uri(httpretty.GET,
                               re.compile("http://169.254.169.254/(.*)"),
                               body=mangle_funct)
    
def teardown_mangle():
    if httpretty.is_enabled():
        httpretty.disable()<commit_msg>Improved error handling for mangle function<commit_after>import logging
import re
import sys

import requests
import httpretty

NOT_FOUND = '404 - Not Found'
# 'http://ec2-54-254-24-239.ap-southeast-1.compute.amazonaws.com/?url=%s'
VULN_URL = None


def mangle(method, uri, headers):
    '''
    See help at common_arguments.add_mangle_arguments to understand what this is.
    '''
    mangled_url = VULN_URL % uri
    
    logging.debug('Requesting %s' % mangled_url)
    response = requests.get(mangled_url)
    
    code = 200
    if NOT_FOUND in response.text:
        code = 404
        
    return (code, headers, response.text)
    
def setup_mangle(cmd_args):
    if cmd_args.mangle_function is not None:
        # We need to verify that the user configured function exists
        try:
            mangle_module_funct = cmd_args.mangle_function
            mangle_module = '.'.join(mangle_module_funct.split('.')[:-1])
            mangle_funct = mangle_module_funct.split('.')[-1]
            
            _temp = __import__(mangle_module, globals(), locals(), ['mangle_funct'], -1)
            mangle_funct = getattr(_temp, mangle_funct)
        except:
            msg = 'The user configured mangle function %s does not exist.'
            logging.critical(msg % mangle_funct)
            sys.exit(1)
        
        if VULN_URL is None:
            logging.critical('URL mangling to access remote instance meta-data'\
                             ' requires you to code your own mangle function'\
                             ' and configure it.')
            sys.exit(1)

        # enable HTTPretty so that it will monkey patch the socket module
        httpretty.enable()
    
        httpretty.register_uri(httpretty.GET,
                               re.compile("http://169.254.169.254/(.*)"),
                               body=mangle_funct)
    
def teardown_mangle():
    if httpretty.is_enabled():
        httpretty.disable()
{'fraction_non_alphanumeric': 0.09835149483095837, 'fraction_numerical': 0.020676166526962837, 'mean_word_length': 2.126637554585153, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2055948', 'n_tokens_mistral': 1142, 'n_tokens_neox': 1085, 'n_words': 258}