addon/addon.py

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()