1 files changed, 34 insertions, 8 deletions
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()