538 lines
18 KiB
Python
538 lines
18 KiB
Python
import asyncio
|
||
import re
|
||
from typing import Annotated
|
||
|
||
from fastapi import APIRouter, Depends, File, HTTPException, UploadFile, status
|
||
from sqlalchemy import func, select, distinct
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy.orm import selectinload
|
||
|
||
from app.dependencies import get_current_user, get_db
|
||
from app.models.device import Device
|
||
from app.models.object import Object
|
||
from app.models.ssh_key import ObjectSSHCredential, ObjectCategoryCredential
|
||
from app.models.tag import Tag
|
||
from app.models.user import User
|
||
from app.schemas.admin import SheetsSyncRequest
|
||
from app.schemas.object import (
|
||
ObjectCreate,
|
||
ObjectImportError,
|
||
ObjectImportPreview,
|
||
ObjectImportResult,
|
||
ObjectImportRow,
|
||
ObjectImportSkipped,
|
||
ObjectMetaResponse,
|
||
ObjectRead,
|
||
ObjectUpdate,
|
||
SyncResult,
|
||
DiscoveryStarted,
|
||
DiscoveryStatus,
|
||
)
|
||
from app.services import crypto
|
||
from app.services.database import AsyncSessionLocal
|
||
from app.services.device_discovery import discover_via_ssh
|
||
from app.services.object_db import sync_from_object_db
|
||
from app.services.sheets import sync_from_sheets
|
||
from app.services.sse import sse_manager, SSEEvent
|
||
|
||
router = APIRouter(prefix="/api/v1/objects", tags=["objects"])
|
||
|
||
_discovery_tasks: dict[int, asyncio.Task] = {}
|
||
|
||
|
||
def _discovery_task_id(obj_id: int) -> str:
|
||
return f"disc-{obj_id}"
|
||
|
||
|
||
def _is_discovery_running(obj_id: int) -> bool:
|
||
task = _discovery_tasks.get(obj_id)
|
||
return task is not None and not task.done()
|
||
|
||
|
||
def _forget_discovery_task(obj_id: int, task: asyncio.Task) -> None:
|
||
if _discovery_tasks.get(obj_id) is task:
|
||
_discovery_tasks.pop(obj_id, None)
|
||
|
||
|
||
def _apply_credentials(obj: Object, db_password: str | None, ssh_password: str | None) -> None:
|
||
if db_password is not None:
|
||
obj.db_pass_enc = crypto.encrypt(db_password)
|
||
if ssh_password is not None:
|
||
obj.ssh_pass_enc = crypto.encrypt(ssh_password)
|
||
|
||
|
||
async def _load_object(db: AsyncSession, obj_id: int) -> Object:
|
||
result = await db.execute(
|
||
select(Object)
|
||
.options(
|
||
selectinload(Object.tags),
|
||
selectinload(Object.ssh_credentials),
|
||
selectinload(Object.category_credentials),
|
||
)
|
||
.where(Object.id == obj_id)
|
||
)
|
||
obj = result.scalar_one_or_none()
|
||
if obj is None:
|
||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Object not found")
|
||
return obj
|
||
|
||
|
||
async def _apply_ssh_credentials(db: AsyncSession, obj: Object, cred_list: list, cat_list: list | None) -> None:
|
||
"""Replace SSH credentials and category credentials for an object."""
|
||
from app.schemas.ssh_key import ObjectSSHCredentialCreate, ObjectCategoryCredentialCreate
|
||
|
||
# Delete existing and re-create (cascade via relationship)
|
||
for existing in list(obj.ssh_credentials):
|
||
await db.delete(existing)
|
||
for cred_data in cred_list:
|
||
new_cred = ObjectSSHCredential(
|
||
object_id=obj.id,
|
||
type=cred_data.type,
|
||
username=cred_data.username,
|
||
password_enc=crypto.encrypt(cred_data.password) if cred_data.password else None,
|
||
ssh_key_id=cred_data.ssh_key_id,
|
||
priority=cred_data.priority,
|
||
)
|
||
db.add(new_cred)
|
||
|
||
if cat_list is not None:
|
||
for existing in list(obj.category_credentials):
|
||
await db.delete(existing)
|
||
for cat_data in cat_list:
|
||
new_cat = ObjectCategoryCredential(
|
||
object_id=obj.id,
|
||
category=cat_data.category,
|
||
protocol=cat_data.protocol,
|
||
type=cat_data.type,
|
||
username=cat_data.username,
|
||
password_enc=crypto.encrypt(cat_data.password) if cat_data.password else None,
|
||
ssh_key_id=cat_data.ssh_key_id,
|
||
priority=cat_data.priority,
|
||
)
|
||
db.add(new_cat)
|
||
|
||
|
||
@router.get("/meta", response_model=ObjectMetaResponse)
|
||
async def get_objects_meta(
|
||
db: AsyncSession = Depends(get_db),
|
||
_: User = Depends(get_current_user),
|
||
):
|
||
"""Return distinct city and client values for autocomplete."""
|
||
cities_result = await db.execute(
|
||
select(distinct(Object.city)).where(Object.city.isnot(None)).order_by(Object.city)
|
||
)
|
||
clients_result = await db.execute(
|
||
select(distinct(Object.client)).where(Object.client.isnot(None)).order_by(Object.client)
|
||
)
|
||
return ObjectMetaResponse(
|
||
cities=[r for (r,) in cities_result.all()],
|
||
clients=[r for (r,) in clients_result.all()],
|
||
)
|
||
|
||
|
||
@router.get("", response_model=list[ObjectRead])
|
||
async def list_objects(
|
||
is_active: bool | None = None,
|
||
tag_id: int | None = None,
|
||
city: str | None = None,
|
||
client: str | None = None,
|
||
db: AsyncSession = Depends(get_db),
|
||
_: User = Depends(get_current_user),
|
||
):
|
||
query = select(Object).options(
|
||
selectinload(Object.tags),
|
||
selectinload(Object.ssh_credentials),
|
||
selectinload(Object.category_credentials),
|
||
)
|
||
if is_active is not None:
|
||
query = query.where(Object.is_active == is_active)
|
||
if tag_id is not None:
|
||
query = query.where(Object.tags.any(Tag.id == tag_id))
|
||
if city is not None:
|
||
query = query.where(Object.city == city)
|
||
if client is not None:
|
||
query = query.where(Object.client == client)
|
||
query = query.order_by(Object.city, Object.client, Object.name)
|
||
result = await db.execute(query)
|
||
objects = result.scalars().all()
|
||
if objects:
|
||
ids = [o.id for o in objects]
|
||
cnt_rows = await db.execute(
|
||
select(Device.object_id, func.count(Device.id).label("cnt"))
|
||
.where(Device.object_id.in_(ids))
|
||
.group_by(Device.object_id)
|
||
)
|
||
counts = dict(cnt_rows.all())
|
||
for obj in objects:
|
||
obj.device_count = counts.get(obj.id, 0)
|
||
return objects
|
||
|
||
|
||
@router.post("", response_model=ObjectRead, status_code=status.HTTP_201_CREATED)
|
||
async def create_object(
|
||
body: ObjectCreate,
|
||
db: AsyncSession = Depends(get_db),
|
||
_: User = Depends(get_current_user),
|
||
):
|
||
data = body.model_dump(exclude={"db_password", "ssh_password", "tag_ids", "ssh_credentials", "category_credentials"})
|
||
obj = Object(**data)
|
||
_apply_credentials(obj, body.db_password, body.ssh_password)
|
||
|
||
if body.tag_ids:
|
||
tags_result = await db.execute(select(Tag).where(Tag.id.in_(body.tag_ids)))
|
||
obj.tags = list(tags_result.scalars().all())
|
||
|
||
db.add(obj)
|
||
await db.flush()
|
||
|
||
if body.ssh_credentials:
|
||
await _apply_ssh_credentials(db, obj, body.ssh_credentials, body.category_credentials)
|
||
await db.flush()
|
||
|
||
await db.refresh(obj, ["tags", "ssh_credentials", "category_credentials"])
|
||
return obj
|
||
|
||
|
||
@router.get("/{obj_id}", response_model=ObjectRead)
|
||
async def get_object(
|
||
obj_id: int,
|
||
db: AsyncSession = Depends(get_db),
|
||
_: User = Depends(get_current_user),
|
||
):
|
||
return await _load_object(db, obj_id)
|
||
|
||
|
||
@router.patch("/{obj_id}", response_model=ObjectRead)
|
||
async def update_object(
|
||
obj_id: int,
|
||
body: ObjectUpdate,
|
||
db: AsyncSession = Depends(get_db),
|
||
_: User = Depends(get_current_user),
|
||
):
|
||
obj = await _load_object(db, obj_id)
|
||
|
||
data = body.model_dump(exclude_none=True, exclude={"db_password", "ssh_password", "tag_ids", "ssh_credentials", "category_credentials"})
|
||
for field, value in data.items():
|
||
setattr(obj, field, value)
|
||
|
||
_apply_credentials(obj, body.db_password, body.ssh_password)
|
||
|
||
if body.tag_ids is not None:
|
||
tags_result = await db.execute(select(Tag).where(Tag.id.in_(body.tag_ids)))
|
||
obj.tags = list(tags_result.scalars().all())
|
||
|
||
if body.ssh_credentials is not None:
|
||
await _apply_ssh_credentials(db, obj, body.ssh_credentials, body.category_credentials)
|
||
|
||
await db.flush()
|
||
await db.refresh(obj)
|
||
return obj
|
||
|
||
|
||
@router.delete("/{obj_id}", status_code=status.HTTP_204_NO_CONTENT)
|
||
async def delete_object(
|
||
obj_id: int,
|
||
force: bool = False,
|
||
db: AsyncSession = Depends(get_db),
|
||
_: User = Depends(get_current_user),
|
||
):
|
||
obj = await _load_object(db, obj_id)
|
||
|
||
if force:
|
||
devices_result = await db.execute(select(Device).where(Device.object_id == obj_id))
|
||
for device in devices_result.scalars().all():
|
||
await db.delete(device)
|
||
await db.flush()
|
||
else:
|
||
count_result = await db.execute(
|
||
select(Device.id).where(Device.object_id == obj_id).limit(1)
|
||
)
|
||
if count_result.scalar_one_or_none() is not None:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_409_CONFLICT,
|
||
detail="Cannot delete object with existing devices. Remove all devices first.",
|
||
)
|
||
await db.delete(obj)
|
||
|
||
|
||
@router.post("/{obj_id}/sync", response_model=SyncResult)
|
||
async def sync_devices(
|
||
obj_id: int,
|
||
db: AsyncSession = Depends(get_db),
|
||
_: User = Depends(get_current_user),
|
||
):
|
||
obj = await _load_object(db, obj_id)
|
||
result = await sync_from_object_db(db, obj)
|
||
return SyncResult(
|
||
created=result.created,
|
||
updated=result.updated,
|
||
skipped=result.skipped,
|
||
errors=result.errors,
|
||
)
|
||
|
||
|
||
async def _run_discovery_background(obj_id: int, task_id: str) -> None:
|
||
"""Run SSH discovery in background, emitting SSE progress events."""
|
||
# Clear stale replay buffer from any previous run of this task_id,
|
||
# then yield the event loop so the frontend SSE connection can register
|
||
# before we start publishing events.
|
||
sse_manager.clear_replay(task_id)
|
||
await asyncio.sleep(0)
|
||
|
||
async def emit(event_type: str, data: dict) -> None:
|
||
await sse_manager.publish(SSEEvent(type=event_type, task_id=task_id, data=data))
|
||
|
||
try:
|
||
async with AsyncSessionLocal() as session:
|
||
async with session.begin():
|
||
result_row = await session.execute(
|
||
select(Object)
|
||
.options(
|
||
selectinload(Object.tags),
|
||
selectinload(Object.ssh_credentials),
|
||
selectinload(Object.category_credentials),
|
||
)
|
||
.where(Object.id == obj_id)
|
||
)
|
||
obj = result_row.scalar_one_or_none()
|
||
if obj is None:
|
||
await emit("disc_done", {"created": 0, "updated": 0, "skipped": 0, "errors": ["Object not found"]})
|
||
return
|
||
discovery_result = await discover_via_ssh(db=session, obj=obj, progress_cb=emit)
|
||
await emit("disc_done", {
|
||
"created": discovery_result.created,
|
||
"updated": discovery_result.updated,
|
||
"skipped": discovery_result.skipped,
|
||
"errors": discovery_result.errors,
|
||
})
|
||
except Exception as exc:
|
||
await emit("disc_done", {"created": 0, "updated": 0, "skipped": 0, "errors": [str(exc)]})
|
||
|
||
|
||
@router.post("/{obj_id}/discover", response_model=DiscoveryStarted)
|
||
async def discover_devices(
|
||
obj_id: int,
|
||
db: AsyncSession = Depends(get_db),
|
||
_: User = Depends(get_current_user),
|
||
):
|
||
"""SSH into each rack PC, parse caps config, discover network devices."""
|
||
await _load_object(db, obj_id)
|
||
task_id = _discovery_task_id(obj_id)
|
||
if _is_discovery_running(obj_id):
|
||
return DiscoveryStarted(task_id=task_id, running=True, already_running=True)
|
||
|
||
task = asyncio.create_task(_run_discovery_background(obj_id, task_id))
|
||
_discovery_tasks[obj_id] = task
|
||
task.add_done_callback(lambda done_task: _forget_discovery_task(obj_id, done_task))
|
||
return DiscoveryStarted(task_id=task_id, running=True, already_running=False)
|
||
|
||
|
||
@router.get("/{obj_id}/discover/status", response_model=DiscoveryStatus)
|
||
async def discovery_status(
|
||
obj_id: int,
|
||
db: AsyncSession = Depends(get_db),
|
||
_: User = Depends(get_current_user),
|
||
):
|
||
await _load_object(db, obj_id)
|
||
return DiscoveryStatus(task_id=_discovery_task_id(obj_id), running=_is_discovery_running(obj_id))
|
||
|
||
|
||
@router.post("/import/preview", response_model=ObjectImportPreview)
|
||
async def preview_objects_import(
|
||
file: Annotated[UploadFile, File()],
|
||
db: AsyncSession = Depends(get_db),
|
||
_: User = Depends(get_current_user),
|
||
):
|
||
"""Dry-run: parse xlsx and return what would be created/skipped."""
|
||
return await _parse_objects_xlsx(db, file, dry_run=True)
|
||
|
||
|
||
@router.post("/import", response_model=ObjectImportResult, status_code=status.HTTP_201_CREATED)
|
||
async def import_objects(
|
||
file: Annotated[UploadFile, File()],
|
||
db: AsyncSession = Depends(get_db),
|
||
_: User = Depends(get_current_user),
|
||
):
|
||
"""Import objects from xlsx. Deduplication by server_ip (skip + warn on duplicate)."""
|
||
preview = await _parse_objects_xlsx(db, file, dry_run=False)
|
||
return ObjectImportResult(
|
||
created=len(preview.would_create),
|
||
skipped=len(preview.would_skip),
|
||
errors=preview.errors,
|
||
)
|
||
|
||
|
||
def _norm_col(name: str) -> str:
|
||
"""Normalize xlsx header: 'server_ip *' → 'server_ip', 'ssh_user (root)' → 'ssh_user'."""
|
||
name = re.sub(r"\s*\*", "", name)
|
||
name = re.sub(r"\s*\(.*?\)", "", name)
|
||
return name.strip().lower()
|
||
|
||
|
||
async def _parse_objects_xlsx(
|
||
db: AsyncSession, file: UploadFile, dry_run: bool
|
||
) -> ObjectImportPreview:
|
||
import io
|
||
try:
|
||
from openpyxl import load_workbook
|
||
except ImportError:
|
||
raise HTTPException(status_code=500, detail="openpyxl not installed")
|
||
|
||
content = await file.read()
|
||
try:
|
||
wb = load_workbook(io.BytesIO(content), read_only=True, data_only=True)
|
||
except Exception:
|
||
raise HTTPException(status_code=400, detail="Не удалось открыть файл. Убедитесь, что файл в формате .xlsx")
|
||
|
||
# Use first sheet (the data sheet)
|
||
ws = wb.worksheets[0]
|
||
rows = list(ws.iter_rows(values_only=True))
|
||
if not rows:
|
||
raise HTTPException(status_code=400, detail="Файл пустой")
|
||
|
||
# Build column index from header row
|
||
header = [_norm_col(str(c)) if c is not None else "" for c in rows[0]]
|
||
col = {name: idx for idx, name in enumerate(header) if name}
|
||
|
||
required = {"name", "city", "server_ip"}
|
||
missing = required - col.keys()
|
||
if missing:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail=f"Отсутствуют обязательные колонки: {', '.join(sorted(missing))}",
|
||
)
|
||
|
||
def _cell(row: tuple, key: str) -> str | None:
|
||
idx = col.get(key)
|
||
if idx is None or idx >= len(row):
|
||
return None
|
||
val = row[idx]
|
||
if val is None:
|
||
return None
|
||
s = str(val).strip()
|
||
return s if s else None
|
||
|
||
would_create: list[ObjectImportRow] = []
|
||
would_skip: list[ObjectImportSkipped] = []
|
||
errors: list[ObjectImportError] = []
|
||
|
||
# Track IPs seen within this file to catch intra-file duplicates
|
||
seen_ips: dict[str, int] = {}
|
||
|
||
# Pre-load existing server_ips from DB for fast lookup
|
||
existing_ips_result = await db.execute(
|
||
select(Object.server_ip, Object.name).where(Object.server_ip.isnot(None))
|
||
)
|
||
existing_ips: dict[str, str] = {ip: name for ip, name in existing_ips_result.all()}
|
||
|
||
for row_num, row in enumerate(rows[1:], start=2):
|
||
# Skip entirely blank rows
|
||
if all(c is None or str(c).strip() == "" for c in row):
|
||
continue
|
||
|
||
name = _cell(row, "name")
|
||
city = _cell(row, "city")
|
||
server_ip = _cell(row, "server_ip")
|
||
|
||
# Validate required fields
|
||
missing_fields = []
|
||
if not name:
|
||
missing_fields.append("name")
|
||
if not city:
|
||
missing_fields.append("city")
|
||
if not server_ip:
|
||
missing_fields.append("server_ip")
|
||
if missing_fields:
|
||
errors.append(ObjectImportError(
|
||
row=row_num,
|
||
reason=f"Обязательные поля не заполнены: {', '.join(missing_fields)}",
|
||
))
|
||
continue
|
||
|
||
# Deduplicate within file
|
||
if server_ip in seen_ips:
|
||
errors.append(ObjectImportError(
|
||
row=row_num,
|
||
reason=f"Дублирующийся IP {server_ip} в файле (строка {seen_ips[server_ip]})",
|
||
))
|
||
continue
|
||
seen_ips[server_ip] = row_num
|
||
|
||
# Deduplicate against DB
|
||
if server_ip in existing_ips:
|
||
would_skip.append(ObjectImportSkipped(
|
||
row=row_num,
|
||
name=name,
|
||
server_ip=server_ip,
|
||
reason=f"Объект с IP {server_ip} уже существует: «{existing_ips[server_ip]}»",
|
||
))
|
||
continue
|
||
|
||
ssh_user = _cell(row, "ssh_user") or "root"
|
||
ssh_password = _cell(row, "ssh_password") or "123123"
|
||
db_via_tunnel_raw = _cell(row, "db_via_tunnel")
|
||
db_via_tunnel = bool(int(db_via_tunnel_raw)) if db_via_tunnel_raw and db_via_tunnel_raw.isdigit() else False
|
||
|
||
import_row = ObjectImportRow(
|
||
row=row_num,
|
||
name=name,
|
||
city=city,
|
||
server_ip=server_ip,
|
||
client=_cell(row, "client"),
|
||
description=_cell(row, "description"),
|
||
notes=_cell(row, "notes"),
|
||
ssh_user=ssh_user,
|
||
db_via_tunnel=db_via_tunnel,
|
||
)
|
||
would_create.append(import_row)
|
||
|
||
if not dry_run:
|
||
obj = Object(
|
||
name=name,
|
||
city=city,
|
||
server_ip=server_ip,
|
||
client=import_row.client,
|
||
description=import_row.description,
|
||
notes=import_row.notes,
|
||
ssh_user=ssh_user,
|
||
db_via_tunnel=db_via_tunnel,
|
||
is_active=True,
|
||
)
|
||
obj.ssh_pass_enc = crypto.encrypt(ssh_password)
|
||
db.add(obj)
|
||
# update local cache so next rows in this file see the newly added IP
|
||
existing_ips[server_ip] = name
|
||
|
||
if not dry_run:
|
||
await db.flush()
|
||
|
||
return ObjectImportPreview(
|
||
would_create=would_create,
|
||
would_skip=would_skip,
|
||
errors=errors,
|
||
)
|
||
|
||
|
||
@router.post("/{obj_id}/sync/sheets", response_model=SyncResult)
|
||
async def sync_devices_from_sheets(
|
||
obj_id: int,
|
||
body: SheetsSyncRequest,
|
||
db: AsyncSession = Depends(get_db),
|
||
_: User = Depends(get_current_user),
|
||
):
|
||
await _load_object(db, obj_id)
|
||
try:
|
||
result = await sync_from_sheets(
|
||
db,
|
||
object_id=obj_id,
|
||
spreadsheet_id=body.spreadsheet_id,
|
||
sheet_name=body.sheet_name,
|
||
)
|
||
except ValueError as exc:
|
||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc))
|
||
return SyncResult(
|
||
created=result.created,
|
||
updated=result.updated,
|
||
skipped=result.skipped,
|
||
errors=result.errors,
|
||
)
|