58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
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()
|