#!/usr/bin/env python3 """Authenticate, optionally upload, search, and issue a temporary source URL.""" import argparse import json import mimetypes import os import sys import time import urllib.error import urllib.parse import urllib.request import uuid from pathlib import Path class ApiError(RuntimeError): pass class NoRedirect(urllib.request.HTTPRedirectHandler): # Do not forward credentials to a redirected host. def redirect_request(self, req, fp, code, msg, headers, newurl): return None # region client class Client: def __init__(self, base_url, token, timeout=30): parsed = urllib.parse.urlsplit(base_url) if (parsed.scheme not in ("http", "https") or not parsed.netloc or parsed.path not in ("", "/") or parsed.query or parsed.fragment or parsed.username or parsed.password): raise ValueError("TKP_URL must be an origin, for example https://tkp.example.com") self.base_url = base_url.rstrip("/") + "/api/v1" self.token = token self.timeout = timeout self.opener = urllib.request.build_opener(NoRedirect()) def request(self, method, path, *, query=None, payload=None, body=None, content_type=None): url = self.base_url + path if query: url += "?" + urllib.parse.urlencode(query, doseq=True) headers = {"Authorization": f"Bearer {self.token}", "Accept": "application/json"} if payload is not None: body = json.dumps(payload).encode() content_type = "application/json" if content_type: headers["Content-Type"] = content_type request = urllib.request.Request(url, data=body, headers=headers, method=method) try: with self.opener.open(request, timeout=self.timeout) as response: return json.load(response) except urllib.error.HTTPError as error: raw = error.read() try: problem = json.loads(raw) detail = problem.get("detail", problem.get("message", error.reason)) fields = problem.get("errors") if fields: detail += " " + json.dumps(fields) except (ValueError, AttributeError): detail = str(error.reason) raise ApiError(f"HTTP {error.code}: {detail}") from None def libraries(self): cursor = None seen = set() while True: page = self.request("GET", "/libraries", query={"limit": 50, **({"cursor": cursor} if cursor else {})}) yield from page["items"] cursor = page["next_cursor"] if cursor is None: return if cursor in seen: raise ApiError("The server returned a repeated pagination cursor") seen.add(cursor) def upload(self, library_id, path): boundary = "tkp-" + uuid.uuid4().hex filename = path.name.replace('"', "_").replace("\r", "_").replace("\n", "_") mime = mimetypes.guess_type(filename)[0] or "application/octet-stream" body = (f'--{boundary}\r\nContent-Disposition: form-data; name="file"; filename="{filename}"\r\n' f"Content-Type: {mime}\r\n\r\n").encode() + path.read_bytes() + f"\r\n--{boundary}--\r\n".encode() return self.request("POST", f"/libraries/{segment(library_id)}/files", body=body, content_type=f"multipart/form-data; boundary={boundary}") # endregion client def segment(value): return urllib.parse.quote(str(value), safe="") # region polling def wait_until_searchable(client, library_id, resource_id, *, wait_seconds=300, interval=3): deadline = time.monotonic() + wait_seconds while True: remaining = deadline - time.monotonic() if remaining <= 0: raise TimeoutError("Ingestion is still pending. Check the resource in the admin panel; do not re-upload automatically.") # Bound each request by the remaining polling window. original_timeout = client.timeout client.timeout = min(original_timeout, remaining) try: resource = client.request("GET", f"/libraries/{segment(library_id)}/resources/{segment(resource_id)}") finally: client.timeout = original_timeout state = resource["ingestion_status"] if state in ("completed", "completed_with_warnings"): return resource if state in ("partial", "failed"): raise ApiError(f"Ingestion ended with {state}. Ask the operator to inspect the resource.") if state not in ("not_started", "running"): raise ApiError(f"Unknown ingestion state: {state}") time.sleep(min(interval, max(0, deadline - time.monotonic()))) # endregion polling # region workflow def run(client, args): context = client.request("GET", "/auth/check") print(f"Workspace: {context['tenant']['name']}") libraries = list(client.libraries()) for library in libraries: print(f"{library['id']}: {library['name']} ({', '.join(library['access'])})") if not args.library: print("Run again with --library ID and --query TEXT.") return selected = next((item for item in libraries if item["id"] == args.library), None) if selected is None or "read" not in selected["access"]: raise ApiError("Choose a listed library with read access") if args.file: if "resource.create" not in selected["access"]: raise ApiError("Uploading also requires resource.create access") uploaded = client.upload(args.library, args.file) print(f"Resource: {uploaded['id']}; duplicate: {uploaded.get('duplicate', False)}") ready = wait_until_searchable(client, args.library, uploaded["id"], wait_seconds=args.wait_seconds) print(f"Ingestion: {ready['ingestion_status']}") result = client.request("GET", "/search/resources", query={ "query": args.query, "library_ids[]": [args.library], "strategy": "lexical", "limit": 5, }) print(result["notice"]) for hit in result["results"]: resource = hit["resource"] print(json.dumps(resource, ensure_ascii=False, indent=2)) if result["results"]: resource_id = result["results"][0]["resource"]["id"] view = client.request("POST", f"/resources/{segment(resource_id)}/view-urls", payload={"ttl_minutes": 5, "disposition": "inline"}) print(f"First candidate source (expires in {view['expires_in_minutes']} minutes): {view['url']}") else: print("No candidates. Try terms from a processed document or a different permitted library.") # endregion workflow def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--library", help="Opaque library ID from the first run") parser.add_argument("--query", help="Lexical search terms") parser.add_argument("--file", type=Path, help="Optional small supported file to upload (read into memory)") parser.add_argument("--wait-seconds", type=float, default=300, help="Maximum ingestion wait (default: 300)") args = parser.parse_args() if args.library and not args.query: parser.error("--library requires --query") if args.file and not args.library: parser.error("--file requires --library and --query") if args.wait_seconds <= 0: parser.error("--wait-seconds must be positive") token = os.environ.get("TKP_TOKEN") if not token or not os.environ.get("TKP_URL"): parser.error("Set TKP_URL (origin) and TKP_TOKEN") try: run(Client(os.environ["TKP_URL"], token), args) except (ApiError, ValueError, KeyError, OSError, TimeoutError, urllib.error.URLError) as error: print(f"Error: {error}", file=sys.stderr) return 1 return 0 if __name__ == "__main__": sys.exit(main())