MyAgent
31 lines
965 B
Python
31 lines
965 B
Python
from fastapi import FastAPI, Header, HTTPException, Depends
|
|
from sqlalchemy.orm import Session
|
|
import datetime
|
|
import uuid
|
|
from database import SessionLocal, APIKey
|
|
|
|
app = FastAPI()
|
|
|
|
def get_db():
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
@app.get('/get-api-key')
|
|
def get_api_key(db: Session = Depends(get_db)):
|
|
new_key = str(uuid.uuid4())
|
|
expires_at = datetime.datetime.utcnow() + datetime.timedelta(hours=24)
|
|
db_key = APIKey(key=new_key, expires_at=expires_at)
|
|
db.add(db_key)
|
|
db.commit()
|
|
return {'api_key': new_key, 'expires_at': expires_at}
|
|
|
|
@app.get('/hello')
|
|
def read_hello(x_api_key: str = Header(...), db: Session = Depends(get_db)):
|
|
db_key = db.query(APIKey).filter(APIKey.key == x_api_key).first()
|
|
if not db_key or db_key.expires_at < datetime.datetime.utcnow():
|
|
raise HTTPException(status_code=403, detail='Invalid or expired API Key')
|
|
return {'message': 'Hello World'}
|