addressgen/backend/app.py

import json
import os
import sys
from datetime import datetime
from typing import Optional

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if ROOT_DIR not in sys.path:
    sys.path.insert(0, ROOT_DIR)

import p2pk
import p2pkh
import p2sh
import p2tr
import p2wpkh

app = FastAPI(title="AddressGen API", version="0.1.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


class P2PKRequest(BaseModel):
    network: str = "mainnet"
    compressed: bool = False
    save_to_file: bool = False
    filename: Optional[str] = None


class P2PKHRequest(BaseModel):
    network: str = "mainnet"
    compressed: bool = True
    save_to_file: bool = False
    filename: Optional[str] = None


class P2WPKHRequest(BaseModel):
    network: str = "mainnet"
    compressed: bool = True
    save_to_file: bool = False
    filename: Optional[str] = None


class P2TRRequest(BaseModel):
    network: str = "mainnet"
    save_to_file: bool = False
    filename: Optional[str] = None


class P2SHRequest(BaseModel):
    network: str = "mainnet"
    m: int = 2
    n: int = 3
    compressed: bool = True
    sort_pubkeys: bool = True
    save_to_file: bool = False
    filename: Optional[str] = None


class SaveRequest(BaseModel):
    data: dict
    filename: Optional[str] = None


def _save_json(data: dict, filename: Optional[str] = None) -> str:
    wallets_dir = os.path.join(ROOT_DIR, "wallets")
    os.makedirs(wallets_dir, exist_ok=True)

    if filename:
        safe = os.path.basename(filename)
        if not safe.endswith(".json"):
            safe += ".json"
        out_name = safe
    else:
        script_type = data.get("script_type", "wallet")
        network = data.get("network", "net")
        stamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
        out_name = f"{script_type}_{network}_{stamp}.json"

    out_path = os.path.join(wallets_dir, out_name)
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=4)
    return out_path


@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/p2pk")
def create_p2pk(req: P2PKRequest):
    try:
        result = p2pk.generate_p2pk(req.network, req.compressed)
        if req.save_to_file:
            result["saved_path"] = _save_json(result, req.filename)
        return result
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc))


@app.post("/p2pkh")
def create_p2pkh(req: P2PKHRequest):
    try:
        result = p2pkh.generate_legacy_address(req.network, req.compressed)
        if req.save_to_file:
            result["saved_path"] = _save_json(result, req.filename)
        return result
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc))


@app.post("/p2wpkh")
def create_p2wpkh(req: P2WPKHRequest):
    try:
        result = p2wpkh.generate_segwit_address(req.network, req.compressed)
        if req.save_to_file:
            result["saved_path"] = _save_json(result, req.filename)
        return result
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc))


@app.post("/p2tr")
def create_p2tr(req: P2TRRequest):
    try:
        result = p2tr.generate_p2tr_address(req.network)
        if req.save_to_file:
            result["saved_path"] = _save_json(result, req.filename)
        return result
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc))


@app.post("/p2sh")
def create_p2sh(req: P2SHRequest):
    try:
        result = p2sh.generate_p2sh_multisig(
            network=req.network,
            m=req.m,
            n=req.n,
            compressed=req.compressed,
            sort_pubkeys=req.sort_pubkeys,
        )
        if req.save_to_file:
            result["saved_path"] = _save_json(result, req.filename)
        return result
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc))


@app.post("/save")
def save_result(req: SaveRequest):
    try:
        saved_path = _save_json(req.data, req.filename)
        return {"saved_path": saved_path}
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc))


def run():
    import uvicorn

    port = int(os.getenv("ADDRESSGEN_PORT", "8732"))
    uvicorn.run("backend.app:app", host="127.0.0.1", port=port, log_level="info")


if __name__ == "__main__":
    run()