Update main.py

main.py

@@ -1,14 +1,123 @@
-from fastapi import FastAPI
-from app.auth import router as auth_router
-from app.routes import router as api_router
-from app.proxy import router as proxy_router
+import os
+import secrets
+import requests
+from fastapi import FastAPI, Depends, HTTPException, Header
+from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, func
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker, Session
+from pydantic import BaseModel
+from dotenv import load_dotenv
 
-app = FastAPI()
+# Load environment variables from .env file
+load_dotenv()
 
-app.include_router(auth_router, prefix="/auth")
-app.include_router(api_router, prefix="/api")
-app.include_router(proxy_router, prefix="/proxy")
+# Environment variables for MySQL and main API settings
+MYSQL_USER = os.getenv("MYSQL_USER")
+MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD")
+MYSQL_HOST = os.getenv("MYSQL_HOST")
+MYSQL_DB = os.getenv("MYSQL_DB")
+MAIN_API_KEY = os.getenv("MAIN_API_KEY")
+MAIN_API_URL = os.getenv("MAIN_API_URL", "https://api.typegpt.net/v1/chat/completions")
+MODEL_NAME = os.getenv("MODEL_NAME", "Image-Generator")
 
-@app.get("/")
-def root():
-    return {"message": "FastAPI API Key System Running!"}
+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()
+
+# User model
+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)
+    created_at = Column(DateTime(timezone=True), server_default=func.now())
+
+# Create tables
+Base.metadata.create_all(bind=engine)
+
+app = FastAPI(title="API Key Generator & Proxy Service")
+
+# Dependency: Database session
+def get_db():
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
+
+# Utility: Generate a unique API key
+def generate_api_key() -> str:
+    return secrets.token_hex(16)
+
+# Pydantic models
+class UserCreate(BaseModel):
+    username: str
+    password: str
+
+class UserOut(BaseModel):
+    id: int
+    username: str
+    api_key: str
+    is_admin: bool
+
+    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
+
+# --- Endpoints ---
+
+# 1. User Registration (generates a unique API key)
+@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)
+    db.add(new_user)
+    db.commit()
+    db.refresh(new_user)
+    return new_user
+
+# 2. User Panel: Get current user info
+@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
+
+# 4. Endpoint to access the main API (proxying request using main API key)
+class RequestPayload(BaseModel):
+    prompt: str
+
+@app.post("/generate")
+def generate_image(payload: RequestPayload, current_user: User = Depends(get_current_user)):
+    headers = {
+        "Authorization": f"Bearer {MAIN_API_KEY}",
+        "Content-Type": "application/json"
+    }
+    data = {
+        "model": MODEL_NAME,
+        "prompt": payload.prompt
+    }
+    response = requests.post(MAIN_API_URL, json=data, headers=headers)
+    if response.status_code != 200:
+        raise HTTPException(status_code=response.status_code, detail="Error from main API")
+    return response.json()