nuiks-sysmon/main.py

from fastapi import FastAPI
import psutil
import time
import datetime
from pydantic import BaseModel
from sqlalchemy import create_engine, Column, Integer, String, DateTime
from sqlalchemy.orm import sessionmaker, declarative_base
from fastapi import HTTPException
import os
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
import threading

if os.environ.get("DATABASE_URL") == None:
    DATABASE_URL = "sqlite:///./wall_messages.db"
    engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
else:
    DATABASE_URL = os.environ["DATABASE_URL"]
    engine = create_engine(DATABASE_URL)
    
    
    
app = FastAPI()


SessionLocal = sessionmaker(bind=engine)
Base = declarative_base()

class LimitsModel(Base):
    __tablename__ = "limits"
    id = Column(Integer, primary_key=True, index=True)
    cpu_threshold = Column(Integer, default=80)
    ram_threshold = Column(Integer, default=80)
    disk_threshold = Column(Integer, default=80)

class OverageEventModel(Base):
    __tablename__ = "overages"
    id = Column(Integer, primary_key=True, index=True)
    cpu_overage = Column(Integer, default=0)
    ram_overage = Column(Integer, default=0)
    disk_overage = Column(Integer, default=0)
    timestamp = Column(DateTime, default=datetime.datetime.now(), nullable=False)

class OverageEvent(BaseModel):
    id: int
    cpu_overage: int
    ram_overage: int
    disk_overage: int
    timestamp: datetime.datetime

    class Config:
        orm_mode = True

Base.metadata.create_all(bind=engine)


@app.on_event("startup")
def startup_event():
    db = SessionLocal()
    if not db.query(LimitsModel).first():
        limits = LimitsModel(
            cpu_threshold=50,
            ram_threshold=50,
            disk_threshold=80
        )
        db.add(limits)
        db.commit()
    db.close()

def check_and_log_overages():
    while True:
        db = SessionLocal()
        limits = db.query(LimitsModel).first()
        if limits:
            cpu_limit = limits.cpu_threshold
            ram_limit = limits.ram_threshold
            disk_limit = limits.disk_threshold
        print("LIMITS: ",cpu_limit, ram_limit, disk_limit)
        cpu = psutil.cpu_percent(interval=0.5)
        ram = psutil.virtual_memory().percent
        disk = psutil.disk_usage('/').percent

        cpu_over = int(cpu > cpu_limit)
        ram_over = int(ram > ram_limit)
        disk_over = int(disk > disk_limit)
        
        last_event = db.query(OverageEventModel).order_by(OverageEventModel.timestamp.desc()).first()
        should_log = False

        if last_event is None:
            should_log = True
        else:
            if (
                last_event.cpu_overage != cpu_over or
                last_event.ram_overage != ram_over or
                last_event.disk_overage != disk_over
            ):
                should_log = True

        if should_log:
            event = OverageEventModel(
                cpu_overage=cpu_over,
                ram_overage=ram_over,
                disk_overage=disk_over,
                timestamp=datetime.datetime.now()
            )
            db.add(event)
            db.commit()
        db.close()
        time.sleep(60)

threading.Thread(target=check_and_log_overages, daemon=True).start()

start_time = time.time()

class Info(BaseModel):
    uptime: datetime.datetime
    cpu: float
    ram: float
    disk: float
    cpu_threshold: int
    ram_threshold: int
    disk_threshold: int
    cpu_normal: bool
    ram_normal: bool
    disk_normal: bool
    
class Thresholds(BaseModel):
    cpu_threshold: int
    ram_threshold: int
    disk_threshold: int


@app.get("/info", response_model=Info)
def get_info():
    db = SessionLocal()
    limits = db.query(LimitsModel).first()
    if not limits:
        db.close()
        raise HTTPException(status_code=404, detail="Limits not found")
    delta_seconds = int(time.time() - start_time)
    cpu_percent = psutil.cpu_percent(interval=0.5)
    ram_percent = psutil.virtual_memory().percent
    disk_percent = psutil.disk_usage('/').percent
    db.close()
    system_information = Info(
        uptime=datetime.datetime.now() - datetime.timedelta(seconds=delta_seconds),
        cpu=cpu_percent,
        ram=ram_percent,
        disk=disk_percent,
        cpu_threshold=limits.cpu_threshold,
        ram_threshold=limits.ram_threshold,
        disk_threshold=limits.disk_threshold,
        cpu_normal= cpu_percent <= limits.cpu_threshold,
        ram_normal=ram_percent <= limits.cpu_threshold,
        disk_normal=disk_percent <= limits.cpu_threshold,
    )
    print(system_information)
    return system_information

@app.get("/limits", response_model=Thresholds)
def get_limits():
    db = SessionLocal()
    limits = db.query(LimitsModel).first()
    if not limits:
        db.close()
        raise HTTPException(status_code=404, detail="Limits not found")
    db.close()
    return {
        "cpu_threshold": limits.cpu_threshold,
        "ram_threshold": limits.ram_threshold,
        "disk_threshold": limits.disk_threshold
    }

@app.post("/limits")
def set_limits(thresholds: Thresholds):
    print (thresholds)
    db = SessionLocal()
    limits = db.query(LimitsModel).first()
    cpu_threshold = thresholds.cpu_threshold
    ram_threshold = thresholds.ram_threshold
    disk_threshold = thresholds.disk_threshold
    if limits:
        limits.cpu_threshold = cpu_threshold
        limits.ram_threshold = ram_threshold
        limits.disk_threshold = disk_threshold
        db.commit()
    db.close()
    return {
        "cpu": cpu_threshold,
        "ram": ram_threshold,
        "disk": disk_threshold
    }
@app.get("/overages", response_model=list[OverageEvent])
def get_overages():
    db = SessionLocal()
    overages = db.query(OverageEventModel).order_by(OverageEventModel.timestamp.desc()).all()
    db.close()
    return overages

@app.get("/overages/{overage_id}", response_model=OverageEvent)
def get_overage(overage_id: int):
    db = SessionLocal()
    overage = db.query(OverageEventModel).filter(OverageEventModel.id == overage_id).first()
    db.close()
    if not overage:
        raise HTTPException(status_code=404, detail="Overage event not found")
    return overage

@app.delete("/overages/{overage_id}")
def delete_overage(overage_id: int):
    db = SessionLocal()
    overage = db.query(OverageEventModel).filter(OverageEventModel.id == overage_id).first()
    if not overage:
        db.close()
        raise HTTPException(status_code=404, detail="Overage event not found")
    db.delete(overage)
    db.commit()
    db.close()
    return {"detail": "Overage event deleted"}

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

app.mount("/", StaticFiles(directory="./webroot", html=True), name="static")