addon/addon.py

59 lines
1.6 KiB
Python
Raw Normal View History

2024-07-01 18:53:05 +03:00
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
2024-07-01 18:53:05 +03:00
);
''')
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),
)
2024-07-01 18:53:05 +03:00
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()