diff --git a/database.py b/database.py new file mode 100644 index 0000000..569fc74 --- /dev/null +++ b/database.py @@ -0,0 +1,17 @@ +from sqlalchemy import create_engine, Column, String, DateTime +from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy.orm import sessionmaker +import datetime + +SQLALCHEMY_DATABASE_URL = 'sqlite:///./test.db' + +engine = create_engine(SQLALCHEMY_DATABASE_URL, connect_args={'check_same_thread': False}) +SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) +Base = declarative_base() + +class APIKey(Base): + __tablename__ = 'api_keys' + key = Column(String, primary_key=True, index=True) + expires_at = Column(DateTime) + +Base.metadata.create_all(bind=engine) diff --git a/main.py b/main.py index f5961e6..4306b75 100644 --- a/main.py +++ b/main.py @@ -1,7 +1,30 @@ -from fastapi import FastAPI +from fastapi import FastAPI, Header, HTTPException, Depends +from sqlalchemy.orm import Session +import datetime +import uuid +from database import SessionLocal, APIKey app = FastAPI() -@app.get("/") -def read_root(): - return {"Hello": "World"} +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'}