ipam_control/app/api/v1/prefixes.py

import ipaddress

from fastapi import APIRouter, Depends, HTTPException, Query
from sqlalchemy import String, and_, cast, func, or_, select, update
from sqlalchemy.orm import Session

from app import schemas as s
from app.db import get_db
from app.models import Address, AddressStatus, Device, Organization, Prefix, PrefixStatus, User, Vrf
from app.security import admin_user, current_user
from app.services import (
    MAX_CAPACITY, apply_update, audit, capacity, commit, count, flush, get_or_404, next_free, utilization,
)

router = APIRouter(dependencies=[Depends(current_user)], tags=["prefixes"])
FREE_LISTING_LIMIT = 4096  # «свободные» строки показываем, только если подсеть не больше /20


# -------------------------------------------------------------------- prefixes
def _usage(db: Session, ids: list[int]) -> dict[int, tuple[int, int]]:
    """{prefix_id: (assigned, stored)} с учётом вложенных префиксов того же VRF."""
    if not ids:
        return {}
    p, c = Prefix.__table__.alias("p"), Prefix.__table__.alias("c")
    a = Address.__table__
    rows = db.execute(
        select(
            p.c.id,
            func.count().filter(a.c.status == AddressStatus.assigned),
            func.count(a.c.id),
        )
        .select_from(p.join(c, and_(c.c.vrf_id == p.c.vrf_id, c.c.prefix.op("<<=")(p.c.prefix)))
                     .join(a, a.c.prefix_id == c.c.id))
        .where(p.c.id.in_(ids))
        .group_by(p.c.id)
    ).all()
    return {r[0]: (r[1], r[2]) for r in rows}


def _depths(db: Session, organization_id: int) -> dict[int, int]:
    parents = dict(db.execute(select(Prefix.id, Prefix.parent_id).where(Prefix.organization_id == organization_id)).all())

    def depth(i: int) -> int:
        d = 0
        while parents.get(i) is not None:
            i, d = parents[i], d + 1
        return d

    return {i: depth(i) for i in parents}


def _capacities(db: Session, organization_id: int) -> dict[int, int]:
    """Ёмкость листа — размер подсети; ёмкость родителя — сумма ёмкостей вложенных листьев."""
    rows = db.execute(select(Prefix.id, Prefix.parent_id, Prefix.prefix).where(Prefix.organization_id == organization_id)).all()
    kids: dict[int, list[int]] = {}
    for i, parent, _ in rows:
        if parent is not None:
            kids.setdefault(parent, []).append(i)
    own = {i: capacity(str(cidr)) for i, _, cidr in rows}
    memo: dict[int, int] = {}

    def cap(i: int) -> int:
        if i not in memo:
            memo[i] = sum(cap(k) for k in kids[i]) if i in kids else own[i]
        return memo[i]

    return {i: min(cap(i), MAX_CAPACITY) for i in own}


def _prefix_outs(db: Session, rows: list[Prefix]) -> list[s.PrefixOut]:
    usage = _usage(db, [r.id for r in rows])
    depths: dict[int, int] = {}
    caps: dict[int, int] = {}
    for org_id in {r.organization_id for r in rows}:
        depths.update(_depths(db, org_id))
        caps.update(_capacities(db, org_id))
    out = []
    for r in rows:
        used, stored = usage.get(r.id, (0, 0))
        cap = caps.get(r.id, capacity(str(r.prefix)))
        out.append(s.PrefixOut(
            id=r.id, organization_id=r.organization_id, vrf_id=r.vrf_id, vrf_name=r.vrf.name,
            prefix=str(r.prefix), family=ipaddress.ip_network(str(r.prefix)).version,
            description=r.description, status=r.status, parent_id=r.parent_id, depth=depths.get(r.id, 0),
            is_pool=r.is_pool, note=r.note, used=used, capacity=cap,
            utilization=utilization(used, cap), addresses_count=stored,
        ))
    return out


def attach_to_tree(db: Session, p: Prefix, keep_parent: bool = False, exclude: frozenset[int] = frozenset()) -> None:
    """Вписывает префикс в дерево его VRF: находит самого узкого родителя и забирает под себя
    вложенные префиксы, чей текущий родитель шире (exclude — уже подчинённые ему, при переносе поддерева)."""
    cidr = str(p.prefix)
    if not keep_parent:
        p.parent_id = db.scalar(
            select(Prefix.id)
            .where(Prefix.vrf_id == p.vrf_id, Prefix.id != p.id, Prefix.prefix.op(">>")(cidr))
            .order_by(func.masklen(Prefix.prefix).desc()).limit(1)
        )
    plen = ipaddress.ip_network(cidr).prefixlen
    inner = db.scalars(
        select(Prefix).where(Prefix.vrf_id == p.vrf_id, Prefix.id != p.id, Prefix.prefix.op("<<")(cidr))
    ).all()
    for c in inner:
        if c.id in exclude:
            continue
        cur = db.get(Prefix, c.parent_id) if c.parent_id else None
        if cur is None or ipaddress.ip_network(str(cur.prefix)).prefixlen < plen:
            c.parent_id = p.id


def _subtree(db: Session, root: Prefix) -> list[Prefix]:
    """Префикс и все вложенные по цепочке parent_id."""
    result, frontier = [root], [root.id]
    while frontier:
        kids = db.scalars(select(Prefix).where(Prefix.parent_id.in_(frontier))).all()
        result += kids
        frontier = [k.id for k in kids]
    return result


def _move_to_vrf(db: Session, p: Prefix, vrf_id: int) -> tuple[dict, int]:
    """Переносит префикс с поддеревом в другой VRF той же организации; возвращает (diff, число префиксов)."""
    vrf = get_or_404(db, Vrf, vrf_id, "VRF")
    if vrf.organization_id != p.organization_id:
        raise HTTPException(422, "Целевой VRF принадлежит другой организации")
    subtree = _subtree(db, p)
    taken = db.scalars(select(Prefix.prefix).where(Prefix.vrf_id == vrf.id, Prefix.prefix.in_([str(m.prefix) for m in subtree]))).all()
    if taken:
        raise HTTPException(409, f"В VRF «{vrf.name}» уже есть: {', '.join(str(x) for x in taken)}")
    old_name = p.vrf.name
    for m in subtree:
        m.vrf = vrf
    p.parent_id = None
    db.flush()
    attach_to_tree(db, p, exclude=frozenset(m.id for m in subtree))
    return {"vrf": f"{old_name} → {vrf.name}", "moved": len(subtree)}, len(subtree)


@router.get("/prefixes", response_model=s.Page[s.PrefixOut])
def list_prefixes(
    organization_id: int | None = None, vrf_id: int | None = None, status: PrefixStatus | None = None,
    family: int | None = Query(None, ge=4, le=6), q: str = "",
    limit: int = Query(100, le=1000), offset: int = 0, db: Session = Depends(get_db),
):
    stmt = select(Prefix)
    if organization_id:
        stmt = stmt.where(Prefix.organization_id == organization_id)
    if vrf_id:
        stmt = stmt.where(Prefix.vrf_id == vrf_id)
    if status:
        stmt = stmt.where(Prefix.status == status)
    if family:
        stmt = stmt.where(func.family(Prefix.prefix) == (4 if family == 4 else 6))
    if q:
        stmt = stmt.where(or_(cast(Prefix.prefix, String).ilike(f"%{q.strip()}%"), Prefix.description.ilike(f"%{q.strip()}%")))
    total = count(db, stmt)
    rows = db.scalars(stmt.order_by(Prefix.vrf_id, Prefix.prefix).limit(limit).offset(offset)).all()
    return s.Page(items=_prefix_outs(db, list(rows)), total=total)


@router.get("/prefixes/{id}", response_model=s.PrefixOut)
def get_prefix(id: int, db: Session = Depends(get_db)):
    return _prefix_outs(db, [get_or_404(db, Prefix, id, "Префикс")])[0]


@router.post("/prefixes", response_model=s.PrefixOut, status_code=201)
def create_prefix(body: s.PrefixIn, db: Session = Depends(get_db), user: User = Depends(admin_user)):
    vrf = get_or_404(db, Vrf, body.vrf_id, "VRF")
    if vrf.organization_id != body.organization_id:
        raise HTTPException(422, "VRF принадлежит другой организации")
    get_or_404(db, Organization, body.organization_id, "Организация")
    if db.scalar(select(Prefix.id).where(Prefix.vrf_id == vrf.id, Prefix.prefix == body.prefix)):
        raise HTTPException(409, "Такой префикс уже есть в этом VRF")

    parent_id = body.parent_id
    if parent_id is not None:
        parent = get_or_404(db, Prefix, parent_id, "Родительский префикс")
        if parent.vrf_id != vrf.id or not ipaddress.ip_network(body.prefix).subnet_of(ipaddress.ip_network(str(parent.prefix))):
            raise HTTPException(422, "Родительский префикс должен содержать новый и быть в том же VRF")
    p = Prefix(**{**body.model_dump(), "parent_id": parent_id})
    db.add(p)
    flush(db, "Такой префикс уже есть в этом VRF")
    attach_to_tree(db, p, keep_parent=parent_id is not None)
    audit(db, user, "prefix", p, "created", str(p.prefix), {"vrf": vrf.name})
    commit(db, "Такой префикс уже есть в этом VRF")
    return _prefix_outs(db, [p])[0]


@router.patch("/prefixes/{id}", response_model=s.PrefixOut)
def update_prefix(id: int, body: s.PrefixUpdate, db: Session = Depends(get_db), user: User = Depends(admin_user)):
    p = get_or_404(db, Prefix, id, "Префикс")
    data = body.model_dump(exclude_unset=True, exclude_none=True)
    new_vrf = data.pop("vrf_id", None)
    changed = {k: str(v) for k, v in apply_update(p, data).items()}
    if new_vrf is not None and new_vrf != p.vrf_id:
        changed.update(_move_to_vrf(db, p, new_vrf)[0])
    audit(db, user, "prefix", p, "updated", str(p.prefix), changed)
    commit(db)
    return _prefix_outs(db, [p])[0]


@router.delete("/prefixes/{id}", status_code=204)
def delete_prefix(id: int, force: bool = False, db: Session = Depends(get_db), user: User = Depends(admin_user)):
    p = get_or_404(db, Prefix, id, "Префикс")
    if not force and count(db, select(Address.id).where(Address.prefix_id == id)):
        raise HTTPException(409, "В префиксе есть адреса; удалите их или используйте force=true")
    db.execute(update(Prefix).where(Prefix.parent_id == id).values(parent_id=p.parent_id))
    audit(db, user, "prefix", p, "deleted", str(p.prefix))
    db.delete(p)
    commit(db)


# ------------------------------------------------------------------- addresses
def _addr_out(a: Address, device_name: str | None = None) -> s.AddressOut:
    return s.AddressOut(
        id=a.id, prefix_id=a.prefix_id, address=s.ip_text(a.address), status=a.status.value,
        dns_name=a.dns_name, description=a.description, device_id=a.device_id, device_name=device_name,
        note=a.note, updated_at=a.updated_at,
    )


def _check_device(db: Session, prefix: Prefix, device_id: int | None):
    if device_id is not None:
        d = get_or_404(db, Device, device_id, "Устройство")
        if d.organization_id != prefix.organization_id:
            raise HTTPException(422, "Устройство принадлежит другой организации")


@router.get("/prefixes/{id}/addresses", response_model=s.AddressPage)
def list_addresses(
    id: int, status: str = "", q: str = "", limit: int = Query(100, le=500), offset: int = 0,
    db: Session = Depends(get_db),
):
    """status: assigned | reserved | deprecated | free | пусто (все; свободные подмешиваются для малых подсетей)."""
    p = get_or_404(db, Prefix, id, "Префикс")
    cap = capacity(str(p.prefix))
    counts = dict(db.execute(select(Address.status, func.count()).where(Address.prefix_id == id).group_by(Address.status)).all())
    stored = sum(counts.values())
    summary = s.AddressSummary(
        assigned=counts.get(AddressStatus.assigned, 0), reserved=counts.get(AddressStatus.reserved, 0),
        deprecated=counts.get(AddressStatus.deprecated, 0), free=max(cap - stored, 0), capacity=cap,
    )
    if status and status not in {"free", *(x.value for x in AddressStatus)}:
        raise HTTPException(422, "Неизвестный статус")

    stmt = select(Address, Device.name).outerjoin(Device, Device.id == Address.device_id).where(Address.prefix_id == id)
    if status and status != "free":
        stmt = stmt.where(Address.status == AddressStatus(status))
    if q:
        like = f"%{q.strip()}%"
        stmt = stmt.where(or_(func.host(Address.address).ilike(like), Address.dns_name.ilike(like), Address.description.ilike(like)))
    rows = [_addr_out(a, dn) for a, dn in db.execute(stmt.order_by(Address.address)).all()] if status != "free" else []

    net = ipaddress.ip_network(str(p.prefix))
    want_free = status == "free" or (not status and not q and cap <= FREE_LISTING_LIMIT)
    if want_free:
        used = {ipaddress.ip_address(s.ip_text(a)) for a in db.scalars(select(Address.address).where(Address.prefix_id == id))}
        hosts = net.hosts() if net.version == 4 and net.prefixlen <= 30 else iter(net)
        need = offset + limit if status == "free" else cap
        free = []
        for ip in hosts:
            if ip not in used:
                free.append(s.AddressOut(id=None, prefix_id=id, address=str(ip), status="free"))
                if len(free) >= need:
                    break
        rows = sorted(rows + free, key=lambda r: ipaddress.ip_address(r.address))
    total = summary.free if status == "free" else (len(rows) if want_free or q or status else stored)
    return s.AddressPage(items=rows[offset:offset + limit], total=total, summary=summary)


@router.post("/prefixes/{id}/addresses", response_model=s.AddressOut, status_code=201)
def create_address(id: int, body: s.AddressIn, db: Session = Depends(get_db), user: User = Depends(admin_user)):
    p = get_or_404(db, Prefix, id, "Префикс")
    if ipaddress.ip_address(body.address) not in ipaddress.ip_network(str(p.prefix)):
        raise HTTPException(422, f"Адрес {body.address} не принадлежит префиксу {p.prefix}")
    _check_device(db, p, body.device_id)
    a = Address(prefix_id=id, **body.model_dump())
    db.add(a)
    flush(db, "Адрес уже есть в этом префиксе")
    audit(db, user, "address", a, "assigned" if a.status == AddressStatus.assigned else "created", body.address)
    commit(db, "Адрес уже есть в этом префиксе")
    db.refresh(a)
    return _addr_out(a)


@router.post("/prefixes/{id}/addresses/next", response_model=s.AddressOut, status_code=201)
def allocate_next(
    id: int, body: s.AddressUpdate | None = None, db: Session = Depends(get_db), user: User = Depends(admin_user)
):
    """Автоназначение первого свободного адреса; только для префиксов с флагом is_pool."""
    p = get_or_404(db, Prefix, id, "Префикс")
    if not p.is_pool:
        raise HTTPException(422, "Префикс не является пулом для автоназначения")
    ip = next_free(db, id, str(p.prefix))
    if ip is None:
        raise HTTPException(409, "В префиксе нет свободных адресов")
    data = body.model_dump(exclude_unset=True, exclude_none=True) if body else {}
    _check_device(db, p, data.get("device_id"))
    a = Address(prefix_id=id, address=ip, **data)
    db.add(a)
    flush(db, "Адрес уже занят, повторите запрос")
    audit(db, user, "address", a, "assigned", ip)
    commit(db, "Адрес уже занят, повторите запрос")
    db.refresh(a)
    return _addr_out(a)


@router.patch("/addresses/{id}", response_model=s.AddressOut)
def update_address(id: int, body: s.AddressUpdate, db: Session = Depends(get_db), user: User = Depends(admin_user)):
    a = get_or_404(db, Address, id, "Адрес")
    data = body.model_dump(exclude_unset=True)
    if data.get("dns_name"):
        s.AddressIn(address=s.ip_text(a.address), dns_name=data["dns_name"])  # валидация FQDN
    _check_device(db, db.get(Prefix, a.prefix_id), data.get("device_id"))
    changed = apply_update(a, data)
    audit(db, user, "address", a, "updated", s.ip_text(a.address), {k: str(v) for k, v in changed.items()})
    commit(db)
    db.refresh(a)
    return _addr_out(a)


@router.delete("/addresses/{id}", status_code=204)
def delete_address(id: int, db: Session = Depends(get_db), user: User = Depends(admin_user)):
    a = get_or_404(db, Address, id, "Адрес")
    audit(db, user, "address", a, "deleted", s.ip_text(a.address))
    db.delete(a)
    commit(db)