Логин + SQL инъекции

API/main.py

@@ -5,9 +5,17 @@ app = FastAPI()
 db = utils.DataBase("")
 
 @app.post("/registration")
-def reg(email: str, password: str) :
+def registration(email: str, password: str) :
     answer = db.registration(email, password)
     if answer["code"] == 1 :
         HTTPException(400, "Email is busy.")
+    else :
+        return answer["data"]
+    
+@app.post("/login")
+def login(email: str, password: str) :
+    answer = db.login(email, password)
+    if answer["code"] == 1 :
+        HTTPException(400, "Incorrect email or password.")
     else :
         return answer["data"]

API/utils.py

@@ -2,6 +2,7 @@ import sqlite3
 import pandas
 import datetime
 import hashlib
+import json
 
 class DataBase :
     def __init__(self, path: str) :
@@ -14,14 +15,20 @@ class DataBase :
                                 email TEXT NOT NULL UNIQUE,
                                 pass_hash TEXT NOT NULL,
                                 register TEXT NOT NULL,
-                                urls TEXT DEFAULT ""
+                                urls TEXT DEFAULT "",
+                                payment TEXT DEFAULT ""
                                 )
                         """)
         
         self.con.commit()
 
-    def read_table(self, sql_command: str) -> dict :
+    def read_table(self, sql_command: str, params=None) -> dict :
-        return pandas.read_sql(sql_command, self.con).to_json()
+        if params == None :
+            params = ()
+        return json.loads(pandas.read_sql(sql_command, self.con, params=params).to_json())
+    
+    def get_hash(text: str) :
+        return hashlib.sha256(text.encode()).hexdigest()
     
     def registration(self, email: str, passwd: str) :
         """
@@ -30,17 +37,29 @@ class DataBase :
         """
         date = datetime.datetime.now()
         try :
-            self.cur.execute(f"""
+            self.cur.execute("""
                                 INSERT INTO users (email, pass_hash, register)
-                                VALUES ({email} {hashlib.sha256(passwd.encode()).hexdigest()} {date.day}-{date.month}-{date.year}_{date.hour}:{date.minute}:{date.second})
+                                VALUES (?, ?, ?)
-                            """)
+                            """, (email, self.get_hash(passwd), f"{date.day}-{date.month}-{date.year} {date.hour}:{date.minute}:{date.second}"))
+            self.con.commit()
         except :
             return {"code": 1}
-        self.con.commit()
+        data = self.read_table("""
-        data = self.read_table(f"""
                                 SELECT id FROM users
-                                WHERE email = {email}
+                                WHERE email = ?
-                               """)
+                               """, (email))
         return {"code": 0, "data": {"id": data["id"]["0"], "email": email}}
         
-    
+    def login(self, email: str, passwd: str) :
+        """
+        0 - Success
+        1 - Incorrect email or password
+        """
+        data = self.read_table("""
+                                    SELECT id, urls, payment FROM users
+                                    WHERE email = ? AND pass_hash = ?
+                                """, (email, self.get_hash(passwd)))
+        if data["id"] == {} :
+            return {"code": 1}
+        else :
+            return {"code": 0, "data": {"id": data["id"]["0"], "email": email, "urls": data["urls"]["0"], "payment": data["payment"]["0"]}}