Files

323 lines
10 KiB
Python

import csv
import io
import json
from fastapi import APIRouter, Depends, HTTPException, Request, status
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session
from sqlalchemy import desc, asc, or_
from datetime import datetime, timezone, timedelta
from typing import List, Optional
from config import get_settings
from database import get_db
from models import Computer, Heartbeat, Notification, AuditLog
from schemas import (
HeartbeatPayload,
ComputerOut,
ComputerDetailOut,
HeartbeatOut,
NotificationOut,
AuditLogOut,
SettingOut,
SettingsUpdate,
)
from auth import get_current_user, require_admin, log_action
from notifications import check_offline_computers, mark_notification_read
from settings_store import get_setting, set_setting, get_all_settings
router = APIRouter(prefix="/api", tags=["api"])
settings = get_settings()
def get_agent_token(request: Request) -> str:
auth = request.headers.get("Authorization", "")
if not auth.startswith("Bearer "):
raise HTTPException(status_code=401, detail="Missing agent token")
return auth[7:]
@router.post("/heartbeat")
def heartbeat(
payload: HeartbeatPayload,
request: Request,
db: Session = Depends(get_db),
):
token = get_agent_token(request)
expected = get_setting(db, "AGENT_TOKEN") or settings.AGENT_TOKEN
if token != expected:
raise HTTPException(status_code=401, detail="Invalid agent token")
now = datetime.now(timezone.utc)
computer = db.query(Computer).filter(Computer.hostname == payload.hostname).first()
if not computer:
computer = Computer(
hostname=payload.hostname,
os_info=payload.os_info,
first_seen=now,
)
db.add(computer)
db.commit()
db.refresh(computer)
computer.last_seen = now
computer.status = "online"
computer.current_user = payload.username
computer.current_ip = payload.local_ip
computer.current_uptime_sec = payload.uptime_sec
computer.current_cpu_percent = payload.cpu_percent
computer.current_ram_percent = payload.ram_percent
computer.current_ram_total_gb = payload.ram_total_gb
computer.current_ram_used_gb = payload.ram_used_gb
computer.current_ram_available_gb = payload.ram_available_gb
computer.current_swap_total_gb = payload.swap_total_gb
computer.current_swap_used_gb = payload.swap_used_gb
computer.current_swap_free_gb = payload.swap_free_gb
computer.current_swap_percent = payload.swap_percent
sessions = payload.sessions
if sessions:
sessions = [s.model_dump() for s in sessions]
heartbeat = Heartbeat(
computer_id=computer.id,
timestamp=now,
username=payload.username,
local_ip=payload.local_ip,
public_ip=payload.public_ip,
uptime_sec=payload.uptime_sec,
cpu_percent=payload.cpu_percent,
ram_percent=payload.ram_percent,
ram_total_gb=payload.ram_total_gb,
ram_used_gb=payload.ram_used_gb,
ram_available_gb=payload.ram_available_gb,
swap_total_gb=payload.swap_total_gb,
swap_used_gb=payload.swap_used_gb,
swap_free_gb=payload.swap_free_gb,
swap_percent=payload.swap_percent,
disk_info=payload.disk_info or [],
cpu_temp=payload.cpu_temp,
load_avg=payload.load_avg,
processes=payload.processes or [],
sessions=sessions,
)
db.add(heartbeat)
db.commit()
return {"ok": True, "computer_id": computer.id}
@router.get("/computers", response_model=List[ComputerOut])
def list_computers(
status: Optional[str] = None,
user: Optional[str] = None,
search: Optional[str] = None,
sort: str = "hostname",
order: str = "asc",
db: Session = Depends(get_db),
current_user=Depends(get_current_user),
):
# Mark old computers as offline and create notifications
check_offline_computers(db)
query = db.query(Computer)
if status and status in ("online", "offline"):
query = query.filter(Computer.status == status)
if user:
query = query.filter(Computer.current_user.ilike(f"%{user}%"))
if search:
query = query.filter(
or_(
Computer.hostname.ilike(f"%{search}%"),
Computer.os_info.ilike(f"%{search}%"),
Computer.current_ip.ilike(f"%{search}%"),
)
)
allowed_sort = {"hostname", "last_seen", "current_user", "status", "current_cpu_percent", "current_ram_percent"}
if sort in allowed_sort:
column = getattr(Computer, sort)
query = query.order_by(asc(column) if order == "asc" else desc(column))
else:
query = query.order_by(asc(Computer.hostname))
return query.all()
@router.get("/computers/{computer_id}", response_model=ComputerDetailOut)
def get_computer(
computer_id: int,
hours: int = 24,
db: Session = Depends(get_db),
user=Depends(get_current_user),
):
computer = db.query(Computer).filter(Computer.id == computer_id).first()
if not computer:
raise HTTPException(status_code=404, detail="Computer not found")
since = datetime.now(timezone.utc) - timedelta(hours=hours)
heartbeats = (
db.query(Heartbeat)
.filter(Heartbeat.computer_id == computer_id, Heartbeat.timestamp >= since)
.order_by(Heartbeat.timestamp)
.all()
)
computer.heartbeats = heartbeats
return computer
@router.delete("/computers/{computer_id}")
def delete_computer(
computer_id: int,
db: Session = Depends(get_db),
admin=Depends(require_admin),
):
computer = db.query(Computer).filter(Computer.id == computer_id).first()
if not computer:
raise HTTPException(status_code=404, detail="Computer not found")
hostname = computer.hostname
db.delete(computer)
db.commit()
log_action(db, admin, "computer_deleted", f"Deleted computer {hostname} (id={computer_id})")
return {"ok": True}
@router.get("/computers/export/csv")
def export_csv(
status: Optional[str] = None,
user: Optional[str] = None,
search: Optional[str] = None,
db: Session = Depends(get_db),
current_user=Depends(get_current_user),
):
query = db.query(Computer)
if status and status in ("online", "offline"):
query = query.filter(Computer.status == status)
if user:
query = query.filter(Computer.current_user.ilike(f"%{user}%"))
if search:
query = query.filter(
or_(
Computer.hostname.ilike(f"%{search}%"),
Computer.os_info.ilike(f"%{search}%"),
Computer.current_ip.ilike(f"%{search}%"),
)
)
computers = query.order_by(Computer.hostname).all()
output = io.StringIO()
writer = csv.writer(output)
writer.writerow([
"ID", "Hostname", "User", "IP", "Status", "OS", "CPU %", "RAM %",
"RAM Used GB", "RAM Total GB", "RAM Available GB", "Swap %",
"Swap Used GB", "Swap Total GB", "Swap Free GB", "Uptime sec",
"Last seen", "First seen",
])
for c in computers:
writer.writerow([
c.id,
c.hostname,
c.current_user or "",
c.current_ip or "",
c.status,
c.os_info or "",
"" if c.current_cpu_percent is None else c.current_cpu_percent,
"" if c.current_ram_percent is None else c.current_ram_percent,
"" if c.current_ram_used_gb is None else c.current_ram_used_gb,
"" if c.current_ram_total_gb is None else c.current_ram_total_gb,
"" if c.current_ram_available_gb is None else c.current_ram_available_gb,
"" if c.current_swap_percent is None else c.current_swap_percent,
"" if c.current_swap_used_gb is None else c.current_swap_used_gb,
"" if c.current_swap_total_gb is None else c.current_swap_total_gb,
"" if c.current_swap_free_gb is None else c.current_swap_free_gb,
"" if c.current_uptime_sec is None else c.current_uptime_sec,
c.last_seen.isoformat() if c.last_seen else "",
c.first_seen.isoformat() if c.first_seen else "",
])
output.seek(0)
filename = f"computers-{datetime.now(timezone.utc).strftime('%Y%m%d-%H%M%S')}.csv"
return StreamingResponse(
io.BytesIO(output.getvalue().encode("utf-8-sig")),
media_type="text/csv; charset=utf-8",
headers={"Content-Disposition": f"attachment; filename={filename}"},
)
@router.get("/notifications", response_model=List[NotificationOut])
def list_notifications(
unread_only: bool = False,
limit: int = 50,
db: Session = Depends(get_db),
user=Depends(get_current_user),
):
query = db.query(Notification)
if unread_only:
query = query.filter(Notification.is_read == False)
return query.order_by(desc(Notification.sent_at)).limit(limit).all()
@router.post("/notifications/{notification_id}/read")
def read_notification(
notification_id: int,
db: Session = Depends(get_db),
user=Depends(get_current_user),
):
if mark_notification_read(db, notification_id):
return {"ok": True}
raise HTTPException(status_code=404, detail="Notification not found")
@router.post("/notifications/read-all")
def read_all_notifications(
db: Session = Depends(get_db),
user=Depends(get_current_user),
):
db.query(Notification).filter(Notification.is_read == False).update(
{"is_read": True, "read_at": datetime.now(timezone.utc)}
)
db.commit()
return {"ok": True}
# Admin: audit log
@router.get("/audit", response_model=List[AuditLogOut])
def list_audit(
limit: int = 200,
db: Session = Depends(get_db),
admin=Depends(require_admin),
):
return db.query(AuditLog).order_by(desc(AuditLog.timestamp)).limit(limit).all()
# Admin: settings
@router.get("/settings")
def list_settings(
db: Session = Depends(get_db),
admin=Depends(require_admin),
):
return get_all_settings(db)
@router.post("/settings")
def update_settings(
payload: SettingsUpdate,
db: Session = Depends(get_db),
admin=Depends(require_admin),
):
changed = []
data = payload.model_dump(exclude_unset=True)
for key, value in data.items():
set_setting(db, key, value)
changed.append(key)
if changed:
log_action(db, admin, "settings_updated", f"Updated: {', '.join(changed)}")
return get_all_settings(db)