third assignment of fastapi

fastapi/auth.py

+from fastapi import Depends, HTTPException
+from jose import JWTError, jwt
+from passlib.context import CryptContext
+from datetime import datetime, timedelta
+from sqlalchemy.orm import Session
+import models, schemas
+from database import SessionLocal
+
+SECRET_KEY = "secret"
+ALGORITHM = "HS256"
+ACCESS_TOKEN_EXPIRE_MINUTES = 30
+
+pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
+
+def get_db():
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
+
+def verify_password(plain_password, hashed_password):
+    return pwd_context.verify(plain_password, hashed_password)
+
+def get_password_hash(password):
+    return pwd_context.hash(password)
+
+def authenticate_user(db, username: str, password: str):
+    user = db.query(models.User).filter(models.User.username == username).first()
+    if not user or not verify_password(password, user.hashed_password):
+        return False
+    return user
+
+def create_access_token(data: dict):
+    to_encode = data.copy()
+    expire = datetime.utcnow() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
+    to_encode.update({"exp": expire})
+    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+
+def get_current_user(token: str = Depends(lambda: "")):
+    credentials_exception = HTTPException(status_code=401, detail="Invalid token")
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        username: str = payload.get("sub")
+        if username is None:
+            raise credentials_exception
+        return username
+    except JWTError:
+        raise credentials_exception
+

fastapi/database.py

+from sqlalchemy import create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+DATABASE_URL = "sqlite:///./todos.db"
+
+engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+Base = declarative_base()

fastapi/main.py

+from fastapi import FastAPI, Depends, WebSocket, WebSocketDisconnect
+from sqlalchemy.orm import Session
+from models import Base, User, Todo
+from database import engine
+from schemas import *
+from auth import *
+from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
+
+app = FastAPI()
+Base.metadata.create_all(bind=engine)
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+active_connections = []
+
+@app.post("/register", response_model=UserOut)
+def register(user: UserCreate, db: Session = Depends(get_db)):
+    hashed_password = get_password_hash(user.password)
+    db_user = User(username=user.username, hashed_password=hashed_password)
+    db.add(db_user)
+    db.commit()
+    db.refresh(db_user)
+    return db_user
+
+@app.post("/token")
+def login(form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_db)):
+    user = authenticate_user(db, form_data.username, form_data.password)
+    if not user:
+        raise HTTPException(status_code=400, detail="Incorrect credentials")
+    token = create_access_token(data={"sub": user.username})
+    return {"access_token": token, "token_type": "bearer"}
+
+@app.post("/todos", response_model=TodoOut)
+def create_todo(todo: TodoCreate, db: Session = Depends(get_db), token: str = Depends(oauth2_scheme)):
+    username = get_current_user(token)
+    user = db.query(User).filter(User.username == username).first()
+    db_todo = Todo(title=todo.title, completed=todo.completed, owner_id=user.id)
+    db.add(db_todo)
+    db.commit()
+    db.refresh(db_todo)
+
+    for connection in active_connections:
+        try:
+            connection.send_json({"event": "new_todo", "todo": db_todo.title})
+        except:
+            pass
+
+    return db_todo
+
+@app.get("/todos", response_model=list[TodoOut])
+def list_todos(db: Session = Depends(get_db), token: str = Depends(oauth2_scheme)):
+    username = get_current_user(token)
+    user = db.query(User).filter(User.username == username).first()
+    return db.query(Todo).filter(Todo.owner_id == user.id).all()
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    active_connections.append(websocket)
+    try:
+        while True:
+            await websocket.receive_text()
+    except WebSocketDisconnect:
+        active_connections.remove(websocket)

fastapi/models.py

+from sqlalchemy import Column, Integer, String, Boolean, ForeignKey
+from sqlalchemy.orm import relationship
+from database import Base
+
+class User(Base):
+    __tablename__ = "users"
+    id = Column(Integer, primary_key=True, index=True)
+    username = Column(String, unique=True, index=True)
+    hashed_password = Column(String)
+
+    todos = relationship("Todo", back_populates="owner")
+
+class Todo(Base):
+    __tablename__ = "todos"
+    id = Column(Integer, primary_key=True, index=True)
+    title = Column(String)
+    completed = Column(Boolean, default=False)
+    owner_id = Column(Integer, ForeignKey("users.id"))
+
+    owner = relationship("User", back_populates="todos")

fastapi/requirement.txt

+fastapi
+uvicorn
+python-jose
+passlib[bcrypt]
+sqlalchemy
+pydantic
+websockets
+

fastapi/schemas.py

+from pydantic import BaseModel
+from typing import Optional
+
+class TodoBase(BaseModel):
+    title: str
+    completed: Optional[bool] = False
+
+class TodoCreate(TodoBase):
+    pass
+
+class TodoOut(TodoBase):
+    id: int
+    class Config:
+        orm_mode = True
+
+class UserCreate(BaseModel):
+    username: str
+    password: str
+
+class UserOut(BaseModel):
+    id: int
+    username: str
+    class Config:
+        orm_mode = True
+