initial commit, logs login attempts to DB

1 files changed, 30 insertions, 0 deletions

diff --git a/db_handler.py b/db_handler.py
new file mode 100644
index 0000000..ded882d
--- /dev/null
+++ b/db_handler.py
@@ -0,0 +1,30 @@
+import os
+
+import psycopg
+import yaml
+
+with open(os.path.join('configs', 'database.yml'), 'r') as file:
+    db_con_params = yaml.safe_load(file.read())
+
+
+def log_connection(ip: str, port: int) -> int:
+    with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
+        with conn.cursor() as cur:
+            cur.execute(
+                "INSERT INTO connection (ip, port)"
+                "VALUES (%(ip)s, %(port)s)"
+                "RETURNING id;",
+                {'ip': ip, 'port': port}
+            )
+            connection_id = cur.fetchone()['id']
+            return connection_id
+
+
+def log_login_attempt(username: str, password: str, connection_id: int):
+    with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
+        with conn.cursor() as cur:
+            cur.execute(
+                "INSERT INTO login_attempt (username, password, connection)"
+                "VALUES (%(username)s, %(password)s, %(connection_id)s);",
+                {'username': username, 'password': password, 'connection_id': connection_id}
+            )