import re import socket import select from unit.main import TestUnit class TestHTTP(TestUnit): def http(self, start_str, **kwargs): sock_type = ( 'ipv4' if 'sock_type' not in kwargs else kwargs['sock_type'] ) port = 7080 if 'port' not in kwargs else kwargs['port'] url = '/' if 'url' not in kwargs else kwargs['url'] http = 'HTTP/1.0' if 'http_10' in kwargs else 'HTTP/1.1' headers = ( {'Host': 'localhost', 'Connection': 'close'} if 'headers' not in kwargs else kwargs['headers'] ) body = b'' if 'body' not in kwargs else kwargs['body'] crlf = '\r\n' if 'addr' not in kwargs: addr = '::1' if sock_type == 'ipv6' else '127.0.0.1' else: addr = kwargs['addr'] sock_types = { 'ipv4': socket.AF_INET, 'ipv6': socket.AF_INET6, 'unix': socket.AF_UNIX, } if 'sock' not in kwargs: sock = socket.socket(sock_types[sock_type], socket.SOCK_STREAM) if ( sock_type == sock_types['ipv4'] or sock_type == sock_types['ipv6'] ): sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) if 'wrapper' in kwargs: sock = kwargs['wrapper'](sock) connect_args = addr if sock_type == 'unix' else (addr, port) try: sock.connect(connect_args) except ConnectionRefusedError: sock.close() return None else: sock = kwargs['sock'] if 'raw' not in kwargs: req = ' '.join([start_str, url, http]) + crlf if body is not b'': if isinstance(body, str): body = body.encode() if 'Content-Length' not in headers: headers['Content-Length'] = len(body) for header, value in headers.items(): if isinstance(value, list): for v in value: req += header + ': ' + str(v) + crlf else: req += header + ': ' + str(value) + crlf req = (req + crlf).encode() + body else: req = start_str sock.sendall(req) if TestUnit.detailed: print('>>>', req, sep='\n') resp = '' if 'no_recv' not in kwargs: enc = 'utf-8' if 'encoding' not in kwargs else kwargs['encoding'] read_timeout = ( 5 if 'read_timeout' not in kwargs else kwargs['read_timeout'] ) resp = self.recvall(sock, read_timeout=read_timeout).decode(enc) if TestUnit.detailed: print('<<<', resp.encode('utf-8'), sep='\n') if 'raw_resp' not in kwargs: resp = self._resp_to_dict(resp) if 'start' not in kwargs: sock.close() return resp return (resp, sock) def delete(self, **kwargs): return self.http('DELETE', **kwargs) def get(self, **kwargs): return self.http('GET', **kwargs) def post(self, **kwargs): return self.http('POST', **kwargs) def put(self, **kwargs): return self.http('PUT', **kwargs) def recvall(self, sock, read_timeout=5, buff_size=4096): data = b'' while select.select([sock], [], [], read_timeout)[0]: try: part = sock.recv(buff_size) except: break data += part if not len(part): break return data def _resp_to_dict(self, resp): m = re.search('(.*?\x0d\x0a?)\x0d\x0a?(.*)', resp, re.M | re.S) if not m: return {} headers_text, body = m.group(1), m.group(2) p = re.compile('(.*?)\x0d\x0a?', re.M | re.S) headers_lines = p.findall(headers_text) status = re.search( '^HTTP\/\d\.\d\s(\d+)|$', headers_lines.pop(0) ).group(1) headers = {} for line in headers_lines: m = re.search('(.*)\:\s(.*)', line) if m.group(1) not in headers: headers[m.group(1)] = m.group(2) elif isinstance(headers[m.group(1)], list): headers[m.group(1)].append(m.group(2)) else: headers[m.group(1)] = [headers[m.group(1)], m.group(2)] return {'status': int(status), 'headers': headers, 'body': body}