diff --git a/functions.py b/functions.py new file mode 100644 index 0000000..ba56b7d --- /dev/null +++ b/functions.py @@ -0,0 +1,16 @@ +from sql import * +from sqlmodel import select + +import hashlib + +def hash_password(password: str, salt: int = None): + password = f"{password}{salt}" + return hashlib.sha256(password.encode()).hexdigest() + +def get_user_by_email(email: str, session) -> User | None: + statement = select(User).where(User.email == email) + return session.exec(statement).first() + +def get_user_by_id(id: str, session: SessionDep) -> User | None: + statement = select(User).where(User.id == id) + return session.exec(statement).first() \ No newline at end of file diff --git a/main.py b/main.py index 294312e..8c1466d 100644 --- a/main.py +++ b/main.py @@ -1,43 +1,20 @@ # Imports -from typing import Annotated +from contextlib import asynccontextmanager +from fastapi import FastAPI, HTTPException +from functions import * +from sql import * -from fastapi import FastAPI, Depends, HTTPException, Query -from sqlmodel import Field, Session, SQLModel, create_engine, select -import hashlib import random -# Initialize the FastAPI App -app = FastAPI() - -# Create the user table - -class User(SQLModel, table=True): - id: int = Field(default=None, primary_key=True) - name: str = Field(index=True) - age: int - email: str = Field(index=True) - password: str - -# SQLModel stuff - -sqlite_file_name = "database.db" -sqlite_url = f"sqlite:///{sqlite_file_name}" - -connect_args = {"check_same_thread": False} -engine = create_engine(sqlite_url, connect_args=connect_args) - -def create_db_and_tables(): - SQLModel.metadata.create_all(engine) - -def get_session(): - with Session(engine) as session: - yield session - -SessionDep = Annotated[Session, Depends(get_session)] - -@app.on_event("startup") -def on_startup(): +# Create DB on startup +# noinspection PyUnusedLocal +@asynccontextmanager +async def lifespan(app: FastAPI): create_db_and_tables() + yield # Code before the yield will run on startup, code after yield won't run until the program is over + +# Initialize the FastAPI App +app = FastAPI(lifespan=lifespan) @app.get("/") def hello_world(): @@ -46,17 +23,15 @@ def hello_world(): @app.get("/hash/{password}") def hash(password: str): salt = random.randint(00000, 99999) - password = password + str(salt) - hashed = hashlib.sha256(password.encode()).hexdigest() + hashed = hash_password(password, salt) return {"hash": hashed, "salt": salt} @app.get("/verify/{password}/{hash}/{salt}") def verify(password: str, hash: str, salt: int): if salt != 0: - password = password + str(salt) - hashed = hashlib.sha256(password.encode()).hexdigest() + hashed = hash_password(password, salt) else: - hashed = hashlib.sha256(password.encode()).hexdigest() + hashed = hash_password(password, salt) if hashed == hash: return {"message": "Password is correct", "correct": True} @@ -65,5 +40,41 @@ def verify(password: str, hash: str, salt: int): @app.get("/hash/no-salt/{password}") def no_salt(password: str): - hashed = hashlib.sha256(password.encode()).hexdigest() - return {"hash": hashed} + hashed = hash_password(password) + return {"hash": hashed, "salt": 0} + +@app.post('/users/create') +async def create_user(user: User, session: SessionDep) -> User | dict[str, str]: + get_user = get_user_by_email(user.email, session) + if get_user is None: + user.password = hash_password(user.password, salt=random.randint(00000, 99999)) + session.add(user) + session.commit() + session.refresh(user) + return user + + return {"message": "User already created"} + +@app.get("/users/{type}") +async def get_user(type: str, session: SessionDep) -> User: + user = get_user_by_id(type, session) + if user is None: + user = get_user_by_email(type, session) + + if not user: + raise HTTPException(status_code=404, detail="User not found") + + return user + +@app.delete("/users/{type}") +async def delete_user(type: str, session: SessionDep) -> User | dict[str, str | bool]: + user = get_user_by_id(type, session) + if user is None: + user = get_user_by_email(type, session) + + if not user: + raise HTTPException(status_code=404, detail="User not found") + + session.delete(user) + session.commit() + return {"message": "User deleted", "completed": True} diff --git a/sql.py b/sql.py new file mode 100644 index 0000000..3774223 --- /dev/null +++ b/sql.py @@ -0,0 +1,27 @@ +from typing import Annotated +from fastapi import Depends +from sqlmodel import Field, Session, SQLModel, create_engine + +# Create the user table +class User(SQLModel, table=True): + id: int = Field(default=None, primary_key=True, index=True) + name: str + age: int + email: str = Field(index=True) + password: str + +# SQLModel stuff +sqlite_file_name = "database.db" +sqlite_url = f"sqlite:///{sqlite_file_name}" + +connect_args = {"check_same_thread": False} +engine = create_engine(sqlite_url, connect_args=connect_args) + +def create_db_and_tables(): + SQLModel.metadata.create_all(engine) + +def get_session(): + with Session(engine) as session: + yield session + +SessionDep = Annotated[Session, Depends(get_session)] \ No newline at end of file