Python Leakix.net Search

leakix_search.py

"""
Search the leakix.net database and save the results to a .jsonl file for further processing.

Requires an API key, get it [here](https://leakix.net/settings/api).
"""

import asyncio
import json
import re
import aiofiles
import aiohttp
import click
from aiopathlib import AsyncPath
from rich.console import Console
from rich.progress import (
    Progress,
    SpinnerColumn,
    TextColumn,
    TimeElapsedColumn,
)

LEAKIX_API_KEY = None
SEARCH_URL = "https://leakix.net/search"


def load_env(path: str = ".env") -> dict[str, str]:
    '''
    Load environment variables from a .env file.

    Args:
        path (str): Path to the .env file. Defaults to ".env".

    Returns:
        dict[str, str]: A dictionary containing the environment variables.

    The .env file should have lines in the format KEY=VALUE. Lines starting
    with # are treated as comments and ignored. Values can be optionally
    enclosed in single or double quotes. If enclosed in double quotes, escape
    sequences will be processed.
      '''
    env: dict[str, str] = {}
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if line.startswith("export "):
                line = line[7:].lstrip()
            key, sep, value = line.partition("=")
            if not sep:
                continue
            key = key.strip()
            if not key:
                continue
            value = value.strip()
            if len(value) >= 2 and value[0] == value[-1] and value[0] in ("'", '"'):
                quote = value[0]
                value = value[1:-1]
                if quote == '"':
                    value = value.encode("utf-8").decode("unicode_escape")
            env[key] = value
    return env

# load the api key from the .env file if not set via env var or command line
LEAKIX_API_KEY = load_env().get("LEAKIX_API_KEY") or LEAKIX_API_KEY


def sanitize_query(q: str) -> str:
    slug = re.sub(r"[^A-Za-z0-9._-]+", "-", q).strip("-")
    return slug or "query"


async def fetch_page(
    session: aiohttp.ClientSession,
    query: str,
    page: int,
    scope: str,
    api_key: str,
    console: Console,
) -> tuple[list | None, int, str | None]:
    headers = {"api-key": api_key, "accept": "application/json"}
    params = {"q": query, "scope": scope, "page": str(page)}
    tries = 0

    while True:
        tries = tries + 1
        console.print(
            f"[cyan]Fetch URL {SEARCH_URL} with params {params}..."
        )
        async with session.get(SEARCH_URL, headers=headers, params=params) as resp:
            if resp.status == 429:
                retry_after_raw = resp.headers.get("Retry-After", "5")
                try:
                    retry_after = max(1, int(float(retry_after_raw)))
                except ValueError:
                    retry_after = 5
                console.print(
                    f"[yellow]429 on page {page}; sleeping {retry_after}s per Retry-After[/]"
                )
                await asyncio.sleep(retry_after)
                continue

            if resp.status >= 400:
                body = await resp.text()
                console.print(
                    f"[green]{resp.status} on page {page}; got {len(body)} bytes of content"
                )
                return None, resp.status, body

            data = await resp.json()
            return data, resp.status, None


async def amain(
    query: str, scope: str, output: str, api_key: str, max_pages: int | None
) -> None:
    console = Console()
    out_path = AsyncPath(output)

    console.print(f"[bold cyan]Query:[/]   {query}")
    console.print(f"[bold cyan]Scope:[/]   {scope}")
    console.print(f"[bold cyan]Output:[/]  {output} (append)")
    if max_pages is not None:
        console.print(f"[bold cyan]Pages:[/]   cap at {max_pages}")

    total_items = 0
    page = 0
    stop_reason = "unknown"

    timeout = aiohttp.ClientTimeout(total=60)
    async with aiohttp.ClientSession(timeout=timeout) as session:
        async with aiofiles.open(out_path, "a") as f:
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                TimeElapsedColumn(),
                console=console,
                transient=False,
            ) as progress:
                task = progress.add_task("starting...", total=None)
                while True:
                    if max_pages is not None and page >= max_pages:
                        stop_reason = f"reached --max-pages={max_pages}"
                        break

                    progress.update(
                        task,
                        description=f"page {page}  |  items so far: {total_items}",
                    )

                    try:
                        data, status, err = await fetch_page(
                            session, query, page, scope, api_key, console
                        )
                    except aiohttp.ClientError as e:
                        stop_reason = f"network error on page {page}: {e!r}"
                        break

                    if err is not None:
                        snippet = (err or "").strip()[:500]
                        console.print(
                            f"[red]HTTP {status} on page {page}[/]\n[red]{snippet}[/]"
                        )
                        stop_reason = f"HTTP {status} on page {page}"
                        break

                    if not data:
                        stop_reason = f"empty page at page {page}"
                        break

                    if not isinstance(data, list):
                        stop_reason = f"unexpected response shape on page {page}: {type(data).__name__}"
                        console.print(f"[red]{stop_reason}[/]")
                        break

                    chunk = "".join(json.dumps(item) + "\n" for item in data)
                    await f.write(chunk)
                    await f.flush()
                    total_items += len(data)
                    console.print(f"[green]Got page {page} with {len(data)} items on it.")
                    page += 1

    console.print()
    console.print(f"[bold green]Done.[/] Got {total_items} items across {page} page(s).")
    console.print(f"[dim]Stop reason: {stop_reason}[/]")
    console.print(f"[dim]Output: {output}[/]")


@click.command()
@click.option(
    "--query",
    "-q",
    default="port:11434",
    show_default=True,
    help="LeakIX search query (e.g. 'port:11434').",
)
@click.option(
    "--scope",
    "-s",
    type=click.Choice(["service", "leak"]),
    default="service",
    show_default=True,
    help="LeakIX search scope.",
)
@click.option(
    "--output",
    "-o",
    default=None,
    help="Output JSONL path. Defaults to results-<sanitized-query>.jsonl",
)
@click.option(
    "--api-key",
    default=LEAKIX_API_KEY,
    envvar="LEAKIX_API_KEY",
    help="LeakIX API key (or set LEAKIX_API_KEY env var).",
)
@click.option(
    "--max-pages",
    type=int,
    default=None,
    help="Optional cap on number of pages to fetch.",
)
def main(
    query: str,
    scope: str,
    output: str | None,
    api_key: str,
    max_pages: int | None,
) -> None:
    if output is None:
        output = f"results-{sanitize_query(query)}.jsonl"
    print("# Hello from leakix-search!\n")
    asyncio.run(amain(query, scope, output, api_key, max_pages))


if __name__ == "__main__":
    main()

requirements.txt

aiofiles
aiopathlib
aiohttp
rich
click