diff --git a/httpz.py b/httpz.py index 28b962e..1542a55 100644 --- a/httpz.py +++ b/httpz.py @@ -27,9 +27,12 @@ except ImportError: exit(1) # ANSI escape codes for colors +BLUE = '\033[34m' +CYAN = '\033[36m' RED = '\033[91m' GREEN = '\033[92m' DARK_GREY = '\033[90m' +YELLOW = '\033[93m' RESET = '\033[0m' # Globals @@ -46,24 +49,50 @@ def vlog(msg: str): logging.info(msg) +def create_session(user_agent: str, timeout: int, proxy: str = None) -> dict: + ''' + Create a custom aiohttp session + + :param user_agent: User agent to use for HTTP requests + :param timeout: Timeout for HTTP requests + ''' + ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE + + headers = {'User-Agent': user_agent} + connector = aiohttp.TCPConnector(ssl=ssl_context) + + session_params = { + 'connector': connector, + 'headers': headers, + 'timeout': aiohttp.ClientTimeout(total=timeout) + } + + return session_params + + def get_dns_servers() -> dict: '''Get a list of DNS servers to use for lookups.''' with urllib.request.urlopen('https://public-dns.info/nameservers.txt') as source: results = source.read().decode().split('\n') + v4_servers = [server for server in results if ':' not in server] v6_servers = [server for server in results if ':' in server] + return {'4': v4_servers, '6': v6_servers} -async def dns_lookup(domain: str, record_type: str, timeout: int) -> list: +async def dns_lookup(domain: str, record_type: str, timeout: int, retry: int) -> list: ''' Resolve DNS information from a domain :param domain: Domain name to resolve :param record_type: DNS record type to resolve :param timeout: Timeout for DNS request + :param retry: Number of times to retry failed requests ''' - for i in range(args.retry): + for i in range(retry): try: version = '4' if record_type == 'A' else '6' if record_type == 'AAAA' else random.choice(['4','6']) nameserver = random.choice(DNS_SERVERS[version]) @@ -82,7 +111,7 @@ async def get_body(source: str, preview: int) -> str: :param source: HTML source of the webpage :param preview: Number of bytes to preview ''' - body_content = re.search(r'