little-things/ia-cdx-search

#!/usr/bin/env python3
import json
import re
import shlex
import sys
import urllib.request


if len(sys.argv) not in (2, 4) or sys.argv[1].lower() in ('-h', '--help') or re.search(r'(^|&)(output|limit|resumekey|showresumekey|page|shownumpages)=', sys.argv[1], re.IGNORECASE):
	print('Usage: ia-cdx-search QUERY [PAGE NUMPAGES]', file = sys.stderr)
	print('Please refer to https://github.com/internetarchive/wayback/tree/master/wayback-cdx-server for the relevant query parameters', file = sys.stderr)
	print('The output, limit, resumeKey, showResumeKey, page, and showNumPages parameters must not be included.', file = sys.stderr)
	print('To resume a search that failed for some reason, provide the page number and number of pages through the second argument instead.', file = sys.stderr)
	print('Output is produces in JSONL format with one line per CDX entry.', file = sys.stderr)
	print('', file = sys.stderr)
	print('Examples:', file = sys.stderr)
	print(" - Subdomains: ia-cdx-search 'url=example.org&collapse=urlkey&fl=original&matchType=domain&filter=original:^https?://[^/]*example\.org(?::[0-9]*)?/'", file = sys.stderr)
	print('   Note that this will only find subdomains whose homepages are in the Wayback Machine. To discover all known subdomains, remove the filter and then extract the domains from the results.', file = sys.stderr)
	print(" - Subdirectories: ia-cdex-search 'url=example.org&collapse=urlkey&fl=original&matchType=domain&filter=original:^https?://[^/]*example\.org(?::[0-9]*)?/[^/]*/'", file = sys.stderr)
	print('   The same caveat applies. The directory must have been retrieved directly without an additional trailing path or query string.', file = sys.stderr)
	sys.exit(1)

query = sys.argv[1]
baseUrl = f'https://web.archive.org/cdx/search/cdx?{query}'
if sys.argv[2:]:
	startPage, numPages = map(int, sys.argv[2:])
else:
	req = urllib.request.Request(f'{baseUrl}&showNumPages=true')
	with urllib.request.urlopen(req) as r:
		if r.getcode() != 200:
			raise RuntimeError(f'Could not fetch number of pages')
		numPages = int(r.read())
	startPage = 0
	print(f'{numPages} pages', file = sys.stderr)

try:
	for page in range(startPage, numPages):
		url = f'{baseUrl}&output=json&page={page}'
		print(f'GET {url}', file = sys.stderr)
		req = urllib.request.Request(url)
		with urllib.request.urlopen(req) as r:
			if r.getcode() != 200:
				raise RuntimeError(f'Could not fetch {url}')
			o = json.load(r)
		assert o, 'got empty response'
		fields = o[0]
		assert all(len(v) == len(fields) for v in o[1:]), 'got unexpected response format'
		for row in o[1:]:
			print(json.dumps(dict(zip(fields, row))))
except (RuntimeError, json.JSONDecodeError, AssertionError):
	print(f'To resume this search from where it crashed, run: ia-cdx-search {shlex.quote(query)} {shlex.quote(page)} {shlex.quote(numPages)}', file = sys.stderr)
	raise