92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
|
||
|
|
import http.server
|
||
|
|
import hashlib
|
||
|
|
import urllib.request
|
||
|
|
import http.server
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
import signal
|
||
|
|
import itertools
|
||
|
|
import requests
|
||
|
|
|
||
|
|
|
||
|
|
class CachingHandler(http.server.BaseHTTPRequestHandler):
|
||
|
|
def _cached_read(self, input, chunksize=4096):
|
||
|
|
# >>> datetime.datetime.utcnow().strftime("%a, %d %b %Y %H:%M:%S GMT")
|
||
|
|
with open(input, 'rb') as data:
|
||
|
|
while True:
|
||
|
|
chunk = data.read(chunksize)
|
||
|
|
if not chunk:
|
||
|
|
break
|
||
|
|
|
||
|
|
yield chunk
|
||
|
|
|
||
|
|
def _saving_read(self, output):
|
||
|
|
with requests.get(self.path, stream=True) as src:
|
||
|
|
if src.status_code != 200:
|
||
|
|
self.send_error(src.status_code)
|
||
|
|
return
|
||
|
|
|
||
|
|
self.send_response(200)
|
||
|
|
self.end_headers()
|
||
|
|
|
||
|
|
try:
|
||
|
|
with open(output, 'wb') as data:
|
||
|
|
for chunk in src.iter_content(chunk_size=4096):
|
||
|
|
if chunk:
|
||
|
|
data.write(chunk)
|
||
|
|
yield(chunk)
|
||
|
|
|
||
|
|
except Exception as err:
|
||
|
|
os.unlink(output)
|
||
|
|
raise err
|
||
|
|
|
||
|
|
def do_GET(self):
|
||
|
|
h = hashlib.sha256()
|
||
|
|
h.update(self.path.encode())
|
||
|
|
path = os.path.join(args.store, h.hexdigest() + ".data")
|
||
|
|
|
||
|
|
chunker = self._saving_read if not os.path.exists(path) else self._cached_read
|
||
|
|
for chunk in chunker(path):
|
||
|
|
self.wfile.write(chunk)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
import argparse
|
||
|
|
import os
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument('--store', type=str, required=True)
|
||
|
|
parser.add_argument('--address', type=str, default="0.0.0.0")
|
||
|
|
parser.add_argument('--port', type=int, default=8000)
|
||
|
|
|
||
|
|
parser.add_argument('cmd', type=str, nargs=1)
|
||
|
|
parser.add_argument('args', type=str, nargs=argparse.REMAINDER)
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
# Instantiate the server before we run the child process
|
||
|
|
os.makedirs(args.store, exist_ok=True)
|
||
|
|
|
||
|
|
server = http.server.HTTPServer(
|
||
|
|
(args.address, args.port),
|
||
|
|
CachingHandler
|
||
|
|
)
|
||
|
|
|
||
|
|
pid = os.fork()
|
||
|
|
if pid == 0:
|
||
|
|
# Start the child process. Tell the parent to shutdown the server when
|
||
|
|
# we are finished.
|
||
|
|
res = subprocess.run(
|
||
|
|
['env', f"http_proxy={args.address}:{args.port}"] + args.cmd + args.args,
|
||
|
|
stderr=sys.stderr, stdout=sys.stdout
|
||
|
|
)
|
||
|
|
|
||
|
|
os.kill(os.getppid(), signal.SIGHUP)
|
||
|
|
sys.exit(res.returncode)
|
||
|
|
else:
|
||
|
|
# Run a server until we get a SIGHUP, then wait for the child to exit.
|
||
|
|
server.serve_forever()
|
||
|
|
(_, status) = os.waitpid(pid)
|
||
|
|
sys.exit(status)
|