nixify repo

2 files changed, 109 insertions, 78 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
diff --git a/src/ssh_node5_net.py b/src/ssh_node5_net.py
index afed9f6..35e9700 100644
--- a/src/ssh_node5_net.py
+++ b/src/ssh_node5_net.py
@@ -1,13 +1,19 @@
 import json
 import logging
 import datetime
+import os
 import random
+import time
 
 import flask
 
 import db_handler
 
 
+db_conninfo = os.environ["DB_CONN"]
+db = db_handler.DBHandler(conninfo=db_conninfo)
+
+
 class CustomFormatter(logging.Formatter):
     grey = "\x1b[90;20m"
     blue = "\x1b[34;20m"
@@ -32,7 +38,12 @@ class CustomFormatter(logging.Formatter):
 
 
 logger = logging.getLogger(__name__)
-logger.root.setLevel(logging.INFO)
+
+
+if os.environ.get('FLASK_DEBUG') == '1':  # Local development
+    logger.root.setLevel(logging.DEBUG)
+else:
+    logger.root.setLevel(logging.INFO)
 
 stream_handler = logging.StreamHandler()
 stream_handler.setFormatter(CustomFormatter())
@@ -43,7 +54,7 @@ app = flask.Flask(__name__, template_folder='templates', static_folder='static',
 
 @app.route("/potm")
 def password_of_the_month():
-    password = db_handler.get_password_of_the_month()
+    password = db.get_password_of_the_month()
     return password
 
 
@@ -52,14 +63,16 @@ def format_time(input: datetime.datetime) -> str:
 
 
 '''
-    histogram = db_handler.get_histogram()
+    histogram = db.get_histogram()
     histogram_data = [a['total_count'] for a in histogram]
     histogram_labels = [format_time(a['date']) for a in histogram]
 '''
 
 min_value = 20
 def get_chart():
-    histogram = db_handler.get_histogram_detailed()
+    start_time = time.time()
+    histogram = db.get_histogram_detailed()
+    db_time = time.time()
     all_dates = sorted(list({d['date'] for d in histogram}))
     by_ip = {}
     for data in histogram:
@@ -71,14 +84,20 @@ def get_chart():
             by_ip[ip] = [0] * len(all_dates)
         by_ip[ip][all_dates.index(data['date'])] += data['total_count']
 
+    histogram_labels = [format_time(a) for a in all_dates]
+
+    processing_time = time.time()
     histogram_chartjs = json.dumps([{
         'label': ip,
         'data': data,
         'backgroundColor': f'hsl({random.randrange(0, 360)}, 50%, 50%)',
         'fill': 'start'
     } for index, (ip, data) in enumerate(by_ip.items())])
+    json_time = time.time()
 
-    histogram_labels = [format_time(a) for a in all_dates]
+    print(f'DB time:         {db_time - start_time}\n'
+          f'Processing time: {processing_time - db_time}\n'
+          f'Json time:       {json_time - processing_time}')
     return histogram_chartjs, histogram_labels
 
 
@@ -94,12 +113,12 @@ def chart_page():
 
 @app.route("/")
 def index():
-    latest_loging_attempts = db_handler.get_latest_login_attempts()
+    latest_loging_attempts = db.get_latest_login_attempts()
     for login_attempt in latest_loging_attempts[0]:
         login_attempt['timestamp'] = format_time(login_attempt['timestamp'])
 
-    top_usernames = db_handler.get_top('username')
-    top_passwords = db_handler.get_top('password')
+    top_usernames = db.get_top('username')
+    top_passwords = db.get_top('password')
 
     histogram_chartjs, histogram_labels = get_chart()
 
@@ -110,3 +129,10 @@ def index():
         top_passwords=top_passwords,
         histogram_data=histogram_chartjs,
         histogram_labels=histogram_labels)
+
+def main():
+    global app
+    app.run()
+
+if __name__ == '__main__':
+    main()