export.py

"""
ChirpStack v4 Device Export Script
────────────────────────────────────────────────────────────────────────────────
Run interactively — prompts for credentials (shared with import_profiles.json).

  python3 export.py
  python3 export.py --base-url https://... --api-key ... --tenant-id ...
"""

import argparse
import csv
import json
import os
import sys
import time
from datetime import datetime, timezone

import requests
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Prompt
from rich.table import Table
from rich import box

console = Console()

LIMIT         = 100
PROFILES_FILE = os.path.join(os.path.dirname(__file__), "import_profiles.json")


# ─── Saved profiles (shared with import.py) ───────────────────────────────────

def load_profiles() -> list[dict]:
    if os.path.exists(PROFILES_FILE):
        with open(PROFILES_FILE) as f:
            return json.load(f)
    return []


def save_profiles(profiles: list[dict]):
    with open(PROFILES_FILE, "w") as f:
        json.dump(profiles, f, indent=2)


def prompt_credentials(args) -> tuple[str, str, str]:
    """Return (base_url, api_key, tenant_id) — from flags or saved profiles."""
    if args.base_url and args.api_key and args.tenant_id:
        return args.base_url.rstrip("/"), args.api_key, args.tenant_id

    profiles = load_profiles()

    if profiles:
        console.print("\n[bold]Saved profiles:[/bold]")
        tbl = Table(box=box.SIMPLE, show_header=True)
        tbl.add_column("#",    justify="right", style="dim", width=3)
        tbl.add_column("Label",  style="cyan")
        tbl.add_column("Tenant ID", style="dim")
        tbl.add_column("API key (preview)", style="dim")
        for i, p in enumerate(profiles, 1):
            tbl.add_row(str(i), p.get("label", "—"), p["tenant_id"], p["api_key"][:16] + "…")
        tbl.add_row("n", "[yellow]Enter new credentials[/yellow]", "", "")
        console.print(tbl)

        choices = [str(i) for i in range(1, len(profiles) + 1)] + ["n"]
        answer  = Prompt.ask("  Select profile", choices=choices, default="n")

        if answer != "n":
            p = profiles[int(answer) - 1]
            base_url = Prompt.ask("  ChirpStack base URL", default="https://console.helium.datamatters.io")
            return base_url.rstrip("/"), p["api_key"], p["tenant_id"]

    console.print()
    base_url  = Prompt.ask("  ChirpStack base URL", default="https://console.helium.datamatters.io")
    label     = Prompt.ask("  Profile label", default="unnamed")
    api_key   = Prompt.ask("  API key")
    tenant_id = Prompt.ask("  Tenant ID")

    profile = {"label": label, "api_key": api_key, "tenant_id": tenant_id}
    if profile not in profiles:
        profiles.append(profile)
        save_profiles(profiles)
        console.print(f"  [dim]Profile '{label}' saved to {PROFILES_FILE}[/dim]")

    return base_url.rstrip("/"), api_key, tenant_id


# ─── API helpers ──────────────────────────────────────────────────────────────

def make_headers(api_key: str) -> dict:
    return {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "Accept": "application/json",
    }


def paginate(base_url: str, headers: dict, url_path: str, result_key: str, params: dict = None):
    offset = 0
    params = params or {}
    while True:
        params.update({"limit": LIMIT, "offset": offset})
        r = requests.get(f"{base_url}{url_path}", headers=headers, params=params)
        r.raise_for_status()
        data   = r.json()
        items  = data.get(result_key, [])
        if not items:
            break
        yield from items
        total  = int(data.get("totalCount", 0))
        offset += len(items)
        if offset >= total:
            break


def get_device(base_url: str, headers: dict, dev_eui: str) -> dict:
    """Fetch full device detail including tags, joinEui and description."""
    r = requests.get(f"{base_url}/api/devices/{dev_eui}", headers=headers)
    if r.status_code == 200:
        body   = r.json()
        detail = body.get("device", {})
        return {
            "tags":        detail.get("tags") or {},
            "joinEui":     detail.get("joinEui", ""),
            "description": detail.get("description", ""),
        }
    return {"tags": {}, "joinEui": "", "description": ""}


def get_device_profile(base_url: str, headers: dict, dp_id: str) -> dict:
    r = requests.get(f"{base_url}/api/device-profiles/{dp_id}", headers=headers)
    if r.status_code == 200:
        return r.json().get("deviceProfile", {})
    return {}


def pick_device_profiles(profiles: list[dict], profile_apps: dict[str, set] = None) -> list[dict]:
    """Show all profiles and let the user select a subset to export."""
    profile_apps = profile_apps or {}
    tbl = Table(box=box.SIMPLE, show_header=True)
    tbl.add_column("#",       justify="right", style="dim", width=3)
    tbl.add_column("Name",    style="cyan")
    tbl.add_column("Region",  style="yellow", width=8)
    tbl.add_column("MAC",     style="dim", max_width=16)
    tbl.add_column("OTAA",    style="green", width=6)
    tbl.add_column("Used by apps", style="dim", max_width=36)
    for i, p in enumerate(profiles, 1):
        apps_str = ", ".join(sorted(profile_apps.get(p["id"], []))) or "–"
        tbl.add_row(
            str(i),
            p.get("name", ""),
            p.get("region", ""),
            p.get("macVersion", ""),
            "✓" if p.get("supportsOtaa") else "–",
            apps_str,
        )
    console.print(tbl)
    console.print(
        "  Select profiles to export (e.g. [cyan]1,3[/cyan]) or press Enter for [bold]all[/bold].\n"
        "  Tip: enter an [cyan]app name substring[/cyan] (e.g. [cyan]SWH[/cyan]) to auto-select all profiles used by matching apps."
    )
    answer = Prompt.ask("  Your choice", default="all").strip()
    if answer.lower() in ("", "all"):
        return profiles
    # numeric selection
    if any(c.isdigit() for c in answer):
        indices = [int(x.strip()) - 1 for x in answer.split(",") if x.strip().isdigit()]
        chosen  = [profiles[i] for i in indices if 0 <= i < len(profiles)]
        return chosen or profiles
    # app name substring filter
    sub = answer.lower()
    chosen = [
        p for p in profiles
        if any(sub in app.lower() for app in profile_apps.get(p["id"], []))
    ]
    if chosen:
        console.print(f"  Auto-selected [green]{len(chosen)}[/green] profile(s) matching [cyan]{answer}[/cyan]")
        return chosen
    console.print("  [yellow]No match — exporting all.[/yellow]")
    return profiles


def get_device_keys(base_url: str, headers: dict, dev_eui: str) -> dict:
    r = requests.get(f"{base_url}/api/devices/{dev_eui}/keys", headers=headers)
    if r.status_code == 200:
        return r.json().get("deviceKeys", {})
    return {}


def get_device_activation(base_url: str, headers: dict, dev_eui: str) -> dict:
    r = requests.get(f"{base_url}/api/devices/{dev_eui}/activation", headers=headers)
    if r.status_code == 200:
        return r.json().get("deviceActivation") or {}
    return {}


# ─── Output helpers ───────────────────────────────────────────────────────────

def slug(base_url: str) -> str:
    """Turn a URL into a short filename-safe slug, e.g. 'helium-datamatters'."""
    host = base_url.replace("https://", "").replace("http://", "").split("/")[0]
    # strip well-known prefixes
    for prefix in ("console.", "chirpstack.", "chirpstack-rest-api."):
        if host.startswith(prefix):
            host = host[len(prefix):]
            break
    # keep first two dot-separated parts and replace dots/colons
    parts = host.split(".")[:2]
    return "-".join(parts).replace(":", "-")


def write_csv(all_devices: list[dict], csv_file: str):
    all_keys = sorted({k for d in all_devices for k in d.keys()})
    with open(csv_file, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(all_keys)
        for d in all_devices:
            row = []
            for key in all_keys:
                value = d.get(key)
                if isinstance(value, list):
                    value = ",".join(map(str, value))
                elif isinstance(value, dict):
                    value = json.dumps(value)
                row.append(value)
            writer.writerow(row)


# ─── Main ─────────────────────────────────────────────────────────────────────

def parse_args():
    p = argparse.ArgumentParser(description="ChirpStack v4 interactive device export tool.")
    p.add_argument("--base-url",   default=None)
    p.add_argument("--api-key",    default=None)
    p.add_argument("--tenant-id",  default=None)
    return p.parse_args()


def main():
    args = parse_args()

    console.print(Panel.fit(
        "[bold cyan]ChirpStack v4 — Device Export Tool[/bold cyan]\n"
        "[dim]Exports all devices + keys to JSON backup and CSV.[/dim]",
        border_style="cyan",
    ))

    base_url, api_key, tenant_id = prompt_credentials(args)
    headers = make_headers(api_key)

    console.print(f"\n[bold]Connecting to:[/bold] {base_url}  [dim](tenant {tenant_id})[/dim]")
    try:
        applications    = list(paginate(base_url, headers, "/api/applications", "result", {"tenantId": tenant_id}))
        dp_list         = list(paginate(base_url, headers, "/api/device-profiles", "result", {"tenantId": tenant_id}))
    except requests.HTTPError as e:
        console.print(f"[red]Connection failed: {e}[/red]")
        sys.exit(1)
    console.print(f"  Found [green]{len(applications)}[/green] application(s), [green]{len(dp_list)}[/green] device profile(s)")

    # ── Fetch lightweight device list to build profile→app map ───────────────
    console.print("  Scanning devices to map profiles to applications...")
    profile_apps: dict[str, set] = {}  # profile_id → set of app names
    with console.status("[bold green]Scanning devices...") as status:
        for app in applications:
            app_id   = app["id"]
            app_name = app.get("name", "")
            for device in paginate(base_url, headers, "/api/devices", "result", {"applicationId": app_id}):
                dp_id = device.get("deviceProfileId", "")
                if dp_id:
                    profile_apps.setdefault(dp_id, set()).add(app_name)

    # ── Device profile export ─────────────────────────────────────────────────
    console.print(f"\n[bold]Device profiles — select which to export:[/bold]")
    selected_dp_list = pick_device_profiles(dp_list, profile_apps)
    all_profiles: list[dict] = []
    with console.status("[bold green]Exporting device profiles...") as status:
        for dp in selected_dp_list:
            status.update(f"[bold green]Fetching profile: {dp['name']}")
            full = get_device_profile(base_url, headers, dp["id"])
            if full:
                all_profiles.append(full)
            time.sleep(0.05)
    console.print(f"  Exported [green]{len(all_profiles)}[/green] device profile(s)")

    all_devices: list[dict] = []

    with console.status("[bold green]Exporting devices...") as status:
        for app in applications:
            app_id   = app["id"]
            app_name = app.get("name", "")
            for device in paginate(base_url, headers, "/api/devices", "result", {"applicationId": app_id}):
                dev_eui = device.get("devEui", "")
                status.update(f"[bold green]{app_name} — {dev_eui} ({len(all_devices)+1})")
                device["applicationName"] = app_name
                device["tenantId"]        = tenant_id

                detail = get_device(base_url, headers, dev_eui)
                device["tags"]        = detail["tags"]
                device["joinEui"]     = detail["joinEui"]
                if detail["description"]:
                    device["description"] = detail["description"]

                keys = get_device_keys(base_url, headers, dev_eui)
                device["appKey"] = keys.get("appKey", "")
                device["nwkKey"] = keys.get("nwkKey", "")

                activation = get_device_activation(base_url, headers, dev_eui)
                device["devAddr"]    = activation.get("devAddr", "")
                device["appSKey"]    = activation.get("appSKey", "")
                device["nwkSEncKey"] = activation.get("nwkSEncKey", "")
                device["sNwkSIntKey"]= activation.get("sNwkSIntKey", "")
                device["fNwkSIntKey"]= activation.get("fNwkSIntKey", "")

                all_devices.append(device)
                time.sleep(0.05)

    console.print(f"  Exported [green]{len(all_devices)}[/green] devices total")

    # ── Summary table ─────────────────────────────────────────────────────────
    app_counts: dict[str, int] = {}
    for d in all_devices:
        app_counts[d["applicationName"]] = app_counts.get(d["applicationName"], 0) + 1
    tbl = Table(title="Applications exported", box=box.ROUNDED, border_style="blue")
    tbl.add_column("Application", style="cyan")
    tbl.add_column("Devices", justify="right", style="green")
    for name, count in sorted(app_counts.items()):
        tbl.add_row(name, str(count))
    console.print(tbl)

    # ── Write files ───────────────────────────────────────────────────────────
    name_slug = slug(base_url)
    ts        = datetime.now(timezone.utc).strftime("%Y-%m-%d_%H%M%S")
    json_file = f"{ts}_{name_slug}_devices_full_backup.json"
    csv_file  = f"{ts}_{name_slug}_devices_full.csv"

    with open(json_file, "w") as f:
        json.dump(all_devices, f, indent=2)
    write_csv(all_devices, csv_file)

    profiles_file = f"{ts}_{name_slug}_device_profiles_backup.json"
    with open(profiles_file, "w") as f:
        json.dump(all_profiles, f, indent=2)

    console.print(f"\n[bold green]✓[/bold green] JSON backup      : [cyan]{json_file}[/cyan]")
    console.print(f"[bold green]✓[/bold green] CSV export       : [cyan]{csv_file}[/cyan]")
    console.print(f"[bold green]✓[/bold green] Profiles backup  : [cyan]{profiles_file}[/cyan]")
    console.print(f"[bold green]✓[/bold green] Total devices    : [green]{len(all_devices)}[/green]")
    console.print(f"[bold green]✓[/bold green] Total profiles   : [green]{len(all_profiles)}[/green]")


if __name__ == "__main__":
    main()