import os import time import logging import sqlite3 from mitmproxy.http import HTTPFlow logging.basicConfig(level=os.getenv('LOG', 'WARN')) storage = os.getenv('STORAGE', 'storage') if not os.path.exists(storage): os.mkdir(storage) db = sqlite3.Connection(os.getenv('SQLITE_DB_PATH', 'archive.db')) with db as conn: conn.executescript(''' create table if not exists data ( id integer primary key, method text not null default "GET", url text not null, code integer not null default 200 ); ''') def response(flow: HTTPFlow) -> None: start = time.time() req = flow.request url = f'{req.scheme}://{req.host}:{req.port}{req.path}' if flow.response is None: logging.warning('response is None: %s', url) return uid: int = 0 with db as conn: code = flow.response.status_code cur = conn.execute( 'insert into data (method, url, code) values (?, ?, ?) returning id', (req.method, url, code), ) uid = cur.fetchone()[0] path = os.path.join(storage, f'{uid}') if not os.path.exists(path): os.mkdir(path) with open(os.path.join(path, 'headers'), 'wb') as fh: logging.debug('%r', flow.response.headers) fh.write(bytes(flow.response.headers)) with open(os.path.join(path, 'body'), 'wb') as fb: body = flow.response.raw_content if body is None: logging.warning('raw_content is None: %s', url) return fb.write(body) logging.info('OK in %.4f sec', time.time() - start) def done() -> None: logging.info('Cleanup') db.close()