Update main.py

main.py

@@ -1,10 +1,11 @@
 import os
 import secrets
+import datetime
 import requests
 from fastapi import FastAPI, Depends, HTTPException, Header, Request
-from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, func
+from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, ForeignKey, func
 from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker, Session
+from sqlalchemy.orm import sessionmaker, Session, relationship
 from pydantic import BaseModel
 from dotenv import load_dotenv
 from fastapi.templating import Jinja2Templates
@@ -13,7 +14,7 @@ from fastapi.staticfiles import StaticFiles
 # Load environment variables from .env file
 load_dotenv()
 
-# Environment variables for MySQL and main API settings
+# Environment variables
 MYSQL_USER = os.getenv("MYSQL_USER")
 MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD")
 MYSQL_HOST = os.getenv("MYSQL_HOST")
@@ -24,32 +25,44 @@ MODEL_NAME = os.getenv("MODEL_NAME", "Image-Generator")
 
 DATABASE_URL = f"mysql+pymysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}/{MYSQL_DB}"
 
-# SQLAlchemy setup
 engine = create_engine(DATABASE_URL)
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 Base = declarative_base()
 
-# Updated User model with credits field
+# ---------------------------
+# Database Models
+# ---------------------------
 class User(Base):
     __tablename__ = "users"
     id = Column(Integer, primary_key=True, index=True)
     username = Column(String(50), unique=True, index=True, nullable=False)
     hashed_password = Column(String(128), nullable=False)
-    api_key = Column(String(64), unique=True, index=True, nullable=False)
     is_admin = Column(Boolean, default=False)
     credits = Column(Integer, default=0)
     created_at = Column(DateTime(timezone=True), server_default=func.now())
+    # Relationship to API keys
+    api_keys = relationship("APIKey", back_populates="owner", cascade="all, delete-orphan")
+
+class APIKey(Base):
+    __tablename__ = "api_keys"
+    id = Column(Integer, primary_key=True, index=True)
+    key = Column(String(64), unique=True, index=True, nullable=False)
+    user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
+    expiry_date = Column(DateTime, nullable=True)
+    active = Column(Boolean, default=True)
+    created_at = Column(DateTime(timezone=True), server_default=func.now())
+    owner = relationship("User", back_populates="api_keys")
 
-# Create tables
 Base.metadata.create_all(bind=engine)
 
-app = FastAPI(title="API Key Platform")
+# ---------------------------
+# FastAPI App & Config
+# ---------------------------
+app = FastAPI(title="Real API Key Platform")
 
-# Mount static files and templates
 app.mount("/static", StaticFiles(directory="static"), name="static")
 templates = Jinja2Templates(directory="templates")
 
-# Dependency: Database session
 def get_db():
     db = SessionLocal()
     try:
@@ -57,11 +70,12 @@ def get_db():
     finally:
         db.close()
 
-# Utility: Generate a unique API key
 def generate_api_key() -> str:
     return secrets.token_hex(16)
 
-# Pydantic models
+# ---------------------------
+# Pydantic Models
+# ---------------------------
 class UserCreate(BaseModel):
     username: str
     password: str
@@ -69,52 +83,91 @@ class UserCreate(BaseModel):
 class UserOut(BaseModel):
     id: int
     username: str
-    api_key: str
     is_admin: bool
     credits: int
+    created_at: datetime.datetime
 
     class Config:
         orm_mode = True
 
-# Dependency: Validate API key from header and return current user
-def get_current_user(x_api_key: str = Header(...), db: Session = Depends(get_db)) -> User:
-    user = db.query(User).filter(User.api_key == x_api_key).first()
-    if not user:
-        raise HTTPException(status_code=401, detail="Invalid API Key")
-    return user
+class APIKeyOut(BaseModel):
+    id: int
+    key: str
+    expiry_date: datetime.datetime | None = None
+    active: bool
+    created_at: datetime.datetime
+
+    class Config:
+        orm_mode = True
+
+class GenerateKeyPayload(BaseModel):
+    expiry_date: datetime.datetime | None = None
+
+class RequestPayload(BaseModel):
+    prompt: str
 
-# --- Endpoints ---
+class CreditPayload(BaseModel):
+    username: str
+    credits: int
+
+class DeactivateKeyPayload(BaseModel):
+    key: str
 
-# 1. User Registration (generates a unique API key)
+# ---------------------------
+# Authentication Dependency
+# ---------------------------
+def get_current_user(x_api_key: str = Header(...), db: Session = Depends(get_db)) -> User:
+    api_key_obj = db.query(APIKey).filter(APIKey.key == x_api_key, APIKey.active == True).first()
+    if not api_key_obj:
+        raise HTTPException(status_code=401, detail="Invalid or inactive API Key")
+    if api_key_obj.expiry_date and datetime.datetime.utcnow() > api_key_obj.expiry_date:
+        raise HTTPException(status_code=401, detail="API Key expired")
+    return api_key_obj.owner
+
+# ---------------------------
+# Endpoints
+# ---------------------------
+
+# Registration: Create a new user and generate a primary API key (no expiry)
 @app.post("/register", response_model=UserOut)
 def register(user: UserCreate, db: Session = Depends(get_db)):
     if db.query(User).filter(User.username == user.username).first():
         raise HTTPException(status_code=400, detail="Username already exists")
-    new_api_key = generate_api_key()
-    # In production, use proper password hashing
-    new_user = User(username=user.username, hashed_password=user.password, api_key=new_api_key, is_admin=False)
+    new_user = User(username=user.username, hashed_password=user.password, is_admin=False)
     db.add(new_user)
     db.commit()
     db.refresh(new_user)
+    primary_key = APIKey(key=generate_api_key(), user_id=new_user.id, expiry_date=None, active=True)
+    db.add(primary_key)
+    db.commit()
+    db.refresh(primary_key)
     return new_user
 
-# 2. User Panel: Get current user info
+# User panel: Get current user details
 @app.get("/user/me", response_model=UserOut)
 def read_user_me(current_user: User = Depends(get_current_user)):
     return current_user
 
-# 3. Admin Panel: List all users (admin-only)
-@app.get("/admin/users", response_model=list[UserOut])
-def list_users(current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
-    if not current_user.is_admin:
-        raise HTTPException(status_code=403, detail="Not authorized")
-    users = db.query(User).all()
-    return users
+# List all API keys belonging to the current user
+@app.get("/user/api_keys", response_model=list[APIKeyOut])
+def get_user_api_keys(current_user: User = Depends(get_current_user)):
+    return current_user.api_keys
 
-# 4. Proxy Endpoint to access the main API
-class RequestPayload(BaseModel):
-    prompt: str
+# Allow user to generate a new API key (with optional expiry date)
+@app.post("/user/generate_key", response_model=APIKeyOut)
+def generate_key(payload: GenerateKeyPayload, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
+    new_key = APIKey(key=generate_api_key(), user_id=current_user.id, expiry_date=payload.expiry_date, active=True)
+    db.add(new_key)
+    db.commit()
+    db.refresh(new_key)
+    return new_key
 
+# Test endpoint for users
+@app.get("/user/test_api")
+def test_api(current_user: User = Depends(get_current_user)):
+    return {"message": "API is working", "username": current_user.username, "credits": current_user.credits}
+
+# Proxy endpoint: Forwards request to main API using the secured main API key
 @app.post("/generate")
 def generate_image(payload: RequestPayload, current_user: User = Depends(get_current_user)):
     headers = {
@@ -130,16 +183,19 @@ def generate_image(payload: RequestPayload, current_user: User = Depends(get_cur
         raise HTTPException(status_code=response.status_code, detail="Error from main API")
     return response.json()
 
-# 5. New endpoint for users to test their API key
-@app.get("/user/test_api")
-def test_api(current_user: User = Depends(get_current_user)):
-    return {"message": "API is working", "username": current_user.username, "credits": current_user.credits}
+# ---------------------------
+# Admin Endpoints
+# ---------------------------
 
-# 6. New endpoint for admin to add credits to a user account
-class CreditPayload(BaseModel):
-    username: str
-    credits: int
+# List all users (admin-only)
+@app.get("/admin/users", response_model=list[UserOut])
+def list_users(current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
+    if not current_user.is_admin:
+        raise HTTPException(status_code=403, detail="Not authorized")
+    users = db.query(User).all()
+    return users
 
+# Add credits to a user's account (admin-only)
 @app.post("/admin/add_credit")
 def add_credit(payload: CreditPayload, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
     if not current_user.is_admin:
@@ -150,16 +206,35 @@ def add_credit(payload: CreditPayload, current_user: User = Depends(get_current_
     user.credits += payload.credits
     db.commit()
     db.refresh(user)
-    return {"message": f"Added {payload.credits} credits to user {user.username}. Total credits: {user.credits}"}
+    return {"message": f"Added {payload.credits} credits to {user.username}. Total credits: {user.credits}"}
 
-# 7. Render Admin Panel UI
-@app.get("/admin/ui")
-def admin_ui(request: Request, current_user: User = Depends(get_current_user)):
+# List all API keys in the system (admin-only)
+@app.get("/admin/api_keys", response_model=list[APIKeyOut])
+def list_all_api_keys(current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
     if not current_user.is_admin:
         raise HTTPException(status_code=403, detail="Not authorized")
+    keys = db.query(APIKey).all()
+    return keys
+
+# Deactivate an API key (admin-only)
+@app.post("/admin/deactivate_key")
+def deactivate_key(payload: DeactivateKeyPayload, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
+    if not current_user.is_admin:
+        raise HTTPException(status_code=403, detail="Not authorized")
+    key_obj = db.query(APIKey).filter(APIKey.key == payload.key).first()
+    if not key_obj:
+        raise HTTPException(status_code=404, detail="API Key not found")
+    key_obj.active = False
+    db.commit()
+    return {"message": f"API Key {payload.key} deactivated."}
+
+# ---------------------------
+# UI Endpoints (Panels)
+# ---------------------------
+@app.get("/admin/ui")
+def admin_ui(request: Request):
     return templates.TemplateResponse("admin.html", {"request": request})
 
-# 8. Render User Panel UI
 @app.get("/user/ui")
-def user_ui(request: Request, current_user: User = Depends(get_current_user)):
+def user_ui(request: Request):
     return templates.TemplateResponse("user.html", {"request": request})