nixify repo

1 files changed, 75 insertions, 70 deletions

diff --git a/src/db_handler.py b/src/db_handler.py
index 9556d6c..e97c8ce 100644
--- a/src/db_handler.py
+++ b/src/db_handler.py
@@ -1,87 +1,92 @@
+import logging
 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())
 
+logger = logging.getLogger(__name__) # Set the logger name, to the name of the module
 
-def get_latest_login_attempts() -> (list[dict], list[str]):
-    with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
-        with conn.cursor() as cur:
-            cur.execute("""
-                SELECT login_attempt.id, username, password, ip, login_attempt.timestamp 
-                FROM login_attempt 
-                JOIN connection on connection.id = login_attempt.connection 
-                ORDER BY login_attempt.id desc limit 20;
-                """)
 
-            login_attempts = cur.fetchall()
-            col_names = [desc[0] for desc in cur.description]
-            return login_attempts, col_names
+class DBHandler:
+    def __init__(self, conninfo: str):
+        self.conninfo = conninfo
+        assert bool(self.get_latest_login_attempts()), "No data from database"
 
+    def get_latest_login_attempts(self) -> (list[dict], list[str]):
+        with psycopg.connect(conninfo=self.conninfo, row_factory=psycopg.rows.dict_row) as conn:
+            with conn.cursor() as cur:
+                cur.execute("""
+                    SELECT login_attempt.id, username, password, ip, login_attempt.timestamp 
+                    FROM login_attempt 
+                    JOIN connection on connection.id = login_attempt.connection 
+                    ORDER BY login_attempt.id desc limit 20;
+                    """)
 
-def get_top(column: str) -> (list[dict], list[str]):
-    if column not in ['username', 'password']:
-        raise ValueError(f'{column} is not allowed')
-    with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
-        with conn.cursor() as cur:
-            cur.execute(psycopg.sql.SQL("""
-                SELECT {column}, COUNT({column}) 
-                FROM login_attempt 
-                GROUP BY {column}
-                ORDER BY COUNT({column}) DESC 
-                LIMIT 20;
-                """).format(column=psycopg.sql.Identifier(column), ))
+                login_attempts = cur.fetchall()
+                col_names = [desc[0] for desc in cur.description]
+                return login_attempts, col_names
 
-            top_usernames = cur.fetchall()
-            col_names = [desc[0] for desc in cur.description]
-            return top_usernames, col_names
 
+    def get_top(self, column: str) -> (list[dict], list[str]):
+        if column not in ['username', 'password']:
+            raise ValueError(f'{column} is not allowed')
+        with psycopg.connect(conninfo=self.conninfo, row_factory=psycopg.rows.dict_row) as conn:
+            with conn.cursor() as cur:
+                cur.execute(psycopg.sql.SQL("""
+                    SELECT {column}, COUNT({column}) 
+                    FROM login_attempt 
+                    GROUP BY {column}
+                    ORDER BY COUNT({column}) DESC 
+                    LIMIT 20;
+                    """).format(column=psycopg.sql.Identifier(column), ))
 
-def get_password_of_the_month() -> str:
-    with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
-        with conn.cursor() as cur:
-            cur.execute("""
-SELECT password
-FROM login_attempt
-WHERE timestamp BETWEEN current_timestamp - interval '1 month' AND timestamp
-GROUP BY password
-ORDER BY COUNT(password) DESC
-LIMIT 1;
-                """)
+                top_usernames = cur.fetchall()
+                col_names = [desc[0] for desc in cur.description]
+                return top_usernames, col_names
 
-            password = cur.fetchone()['password']
-            return password
 
+    def get_password_of_the_month(self) -> str:
+        with psycopg.connect(conninfo=self.conninfo, row_factory=psycopg.rows.dict_row) as conn:
+            with conn.cursor() as cur:
+                cur.execute("""
+    SELECT password
+    FROM login_attempt
+    WHERE timestamp BETWEEN current_timestamp - interval '1 month' AND timestamp
+    GROUP BY password
+    ORDER BY COUNT(password) DESC
+    LIMIT 1;
+                    """)
 
-def get_histogram_detailed() -> str:
-    with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
-        with conn.cursor() as cur:
-            cur.execute("""
-SELECT count(la.id) as total_count, date_trunc('hour', la.timestamp) as date, cn.ip
-FROM login_attempt la
-JOIN connection cn on cn.id = la.connection
-WHERE la.timestamp BETWEEN (select max(timestamp) from login_attempt) - interval '3 days' AND (select max(timestamp) from login_attempt)
-GROUP BY date_trunc('hour', la.timestamp), cn.ip
-ORDER BY COUNT(la.id) DESC
-;
-                """)
-            histogram = cur.fetchall()
-            return histogram
+                password = cur.fetchone()['password']
+                return password
 
 
-def get_histogram_simple() -> str:
-    with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
-        with conn.cursor() as cur:
-            cur.execute("""
-SELECT count(id) as total_count, date_trunc('hour', timestamp) as date
-FROM login_attempt
-GROUP BY date_trunc('hour', timestamp)
-ORDER BY date_trunc('hour', timestamp)
-LIMIT 48
-;
-                """)
-            histogram = cur.fetchall()
-            return histogram
+    def get_histogram_detailed(self) -> str:
+        with psycopg.connect(conninfo=self.conninfo, row_factory=psycopg.rows.dict_row) as conn:
+            with conn.cursor() as cur:
+                cur.execute("""
+    SELECT count(la.id) as total_count, date_trunc('hour', la.timestamp) as date, cn.ip
+    FROM login_attempt la
+    JOIN connection cn on cn.id = la.connection
+    WHERE la.timestamp BETWEEN (select max(timestamp) from login_attempt) - interval '3 days' AND (select max(timestamp) from login_attempt)
+    GROUP BY date_trunc('hour', la.timestamp), cn.ip
+    ORDER BY COUNT(la.id) DESC
+    ;
+                    """)
+                histogram = cur.fetchall()
+                return histogram
+
+
+    def get_histogram_simple(self) -> str:
+        with psycopg.connect(conninfo=self.conninfo, row_factory=psycopg.rows.dict_row) as conn:
+            with conn.cursor() as cur:
+                cur.execute("""
+    SELECT count(id) as total_count, date_trunc('hour', timestamp) as date
+    FROM login_attempt
+    GROUP BY date_trunc('hour', timestamp)
+    ORDER BY date_trunc('hour', timestamp)
+    LIMIT 48
+    ;
+                    """)
+                histogram = cur.fetchall()
+                return histogram