Implement some file changes, updated code & bug fixing

api/api.py

@@ -3,6 +3,7 @@ from contextlib import asynccontextmanager
 from fastapi import FastAPI, HTTPException
 from functions import *
 
+import random
 
 # Create DB on startup
 # noinspection PyUnusedLocal
@@ -17,10 +18,14 @@ app = FastAPI(lifespan=lifespan)
 # Routes
 @app.post("/heroes/create", response_model=HeroPublic)
 def create_hero(hero: HeroCreate, session: SessionDep):
+    salt = random.randint(00000, 99999)
     existing_hero = session.query(Hero).filter(Hero.email == hero.email).first()
     if existing_hero:
         raise HTTPException(status_code=400, detail="Email already registered")
 
+    hero.password = hash_password(hero.password, salt)
+    hero.salt = salt
+
     db_hero = Hero.model_validate(hero)
     session.add(db_hero)
     session.commit()
@@ -28,7 +33,7 @@ def create_hero(hero: HeroCreate, session: SessionDep):
     return db_hero
 
 
-@app.get("/heroes/{id}", response_model=Hero)
+@app.get("/heroes/{id}", response_model=HeroPublic)
 def get_hero(id: int, session: SessionDep):
     hero = session.get(Hero, id)
     if not hero:
@@ -60,14 +65,14 @@ def delete_hero(id: int, session: SessionDep):
     return {"message": "Delete successful", "continue": True}
 
 @app.post("/verify")
-def verify_user(email: str, password: str, session: SessionDep):
-    user = get_hero_by_email(email, session)
+def verify_user(hero: VerifyHero, session: SessionDep):
+    user = get_hero_by_email(hero.email, session)
     if not user:
-        return {"message": "Authentication failed", "continue": False}
+        raise HTTPException(status_code=500, detail="Something went wrong.. Try again later")
     # check if the password is correct
-    authenticate_user = verify_password(password, user.hash, user.salt)
+    authenticate_user = verify_password(hero.password, user.password, user.salt)
 
     if authenticate_user:
-        return {"message": "Authentication successful", "continue": True}
+        return {"message": "Authentication successful", "continue": True, "hero_id": user.id}
 
-    return {"message": "Authentication failed", "continue": False}
+    raise HTTPException(status_code=500, detail="Something went wrong.. Try again later")

api/functions.py

@@ -0,0 +1,31 @@
+# Imports
+from fastapi import HTTPException
+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 verify_password(password: str, hash: str, salt: int) -> bool:
+    hashed_pass = hash_password(password, salt)
+    if hash == hashed_pass:
+        return True
+
+    return False
+
+def get_hero_by_email(email: str, session) -> Hero | None:
+    return session.query(Hero).filter(Hero.email == email).first()
+
+def get_hero_by_id(id: str, session: SessionDep) -> Hero | None:
+    statement = select(Hero).where(Hero.id == id)
+    return session.exec(statement).first()
+
+def get_public_hero(id: int, session: SessionDep) -> HeroPublic:
+    hero = session.get(Hero, id)
+    if not hero:
+        raise HTTPException(status_code=404, detail="Hero not found")
+
+    return hero

api/sql.py

@@ -2,8 +2,7 @@
 from typing import Annotated
 from fastapi import Depends
 from sqlmodel import Field, Session, SQLModel, create_engine
-from sqlalchemy import String
-from sqlalchemy.sql.schema import Column
+from pydantic import BaseModel
 
 # The base hero class
 class HeroBase(SQLModel):
@@ -35,8 +34,13 @@ class HeroUpdate(HeroBase):
     gold: int | None = None
     health: int | None = None
 
+# The class for verifying heroes
+class VerifyHero(BaseModel):
+    email: str
+    password: str
+
 # SQLModel initialization
-sqlite_file_name = "database.db"
+sqlite_file_name = "../database.db"
 sqlite_url = f"sqlite:///{sqlite_file_name}"
 
 connect_args = {"check_same_thread": False}

functions.py

@@ -1,24 +1,33 @@
-# Imports
-from sql import *
-from sqlmodel import select
+from typing import Any
 
-import hashlib
+import requests
 
-def hash_password(password: str, salt: int = None):
-    password = f"{password}{salt}"
-    return hashlib.sha256(password.encode()).hexdigest()
+url="http://127.0.0.1:8000"
 
-def verify_password(password: str, hash: str, salt: int) -> bool:
-    hashed_pass = hash_password(password, salt)
-    if hash == hashed_pass:
-        return True
+def create_account(name: str, email: str, password: str) -> Any | None:
+    data = {
+        "name": name,
+        "email": email,
+        "password": password
+    }
 
-    return False
+    response = requests.post(f"{url}/heroes/create", json=data)
 
-def get_hero_by_email(email: str, session) -> HeroPublic | None:
-    statement = select(Hero.email).where(Hero.email == email)
-    return session.exec(statement).first()
+    if response.status_code == 200:
+        return response.json()
+    else:
+        return False
 
-def get_hero_by_id(id: str, session: SessionDep) -> Hero | None:
-    statement = select(Hero).where(Hero.id == id)
-    return session.exec(statement).first()
+def login(email, password):
+    data = {
+        "email": email,
+        "password": password
+    }
+
+    response = requests.post(f"{url}/verify", json=data)
+
+    if response.status_code == 200:
+        user_response = requests.get(f"{url}/heroes/{response.json()['hero_id']}")
+        return user_response.json()
+    else:
+        return False

main.py

@@ -1,8 +1,9 @@
-import requests
-import json
-import os
+from functions import *
 
-url="http://127.0.0.1:8000"
+import requests
+
+
+user = None
 
 # Does the user have an active account?
 active_account = input("Do you have an account? (y/n): ")
@@ -10,35 +11,28 @@ if active_account == "y":
     email = input("Enter your email: ")
     password = input("Enter your password: ")
 
-    data = {
-        "email": email,
-        "password": password
-    }
+    is_authenticated = login(email, password)
 
-    response = requests.post(f"{url}/verify", json=data)
-
-    if response.status_code == 200:
-        print("Authentication Successful")
-        print(response.json())
+    if is_authenticated is not False:
+        user = is_authenticated
     else:
-        print("Authentication Not Successful")
+        print("Sign in failed. Check your email or password")
 else:
-    create_account = input("Do you want to create an account? (y/n): ")
-    if create_account.lower() == "y":
+    account = input("Do you want to create an account? (y/n): ")
+    if account.lower() == "y":
         name = input("Enter your name: ")
         email = input("Enter your email: ")
         password = input("Enter your password: ")
 
-        data = {
-            "name": name,
-            "email": email,
-            "password": password
-        }
+        is_authenticated = create_account(name, email, password)
 
-        response = requests.post(f"{url}/heroes/create", json=data)
-
-        if response.status_code == 200:
-            print("Authentication Successful")
-            print(response.json())
+        if is_authenticated is not False:
+            user = is_authenticated
         else:
-            print("Authentication Unsuccessful")
+            print("Something went wrong.. Please try again later")
+    else:
+        print("Bye bye..")
+
+while user is not None:
+    print(f"Welcome {user['name']}!")
+    break