Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 from flask import Flask, render_template, request, jsonify, session
 import sqlite3
 from datetime import datetime
@@ -14,73 +15,10 @@ logging.basicConfig(level=logging.DEBUG)
 otp_code = None
 otp_expiration = None
 
-def init_db():
-    with sqlite3.connect('database.db') as conn:
-        cursor = conn.cursor()
-        cursor.execute('''
-            CREATE TABLE IF NOT EXISTS sessions (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                user_id INTEGER,
-                start_time TEXT,
-                end_time TEXT,
-                script_code TEXT,
-                rating INTEGER,
-                comments TEXT,
-                session_type TEXT,
-                duration INTEGER
-            )
-        ''')
-        cursor.execute('''
-            CREATE TABLE IF NOT EXISTS users (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                username TEXT UNIQUE,
-                password TEXT,
-                is_admin INTEGER DEFAULT 0
-            )
-        ''')
-        cursor.execute('''
-            CREATE TABLE IF NOT EXISTS tasks (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                user_id INTEGER,
-                description TEXT,
-                completed INTEGER DEFAULT 0
-            )
-        ''')
-        cursor.execute('''
-            CREATE TABLE IF NOT EXISTS settings (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                user_id INTEGER,
-                setting_name TEXT,
-                setting_value TEXT,
-                UNIQUE(user_id, setting_name)
-            )
-        ''')
-        cursor.execute('''
-            CREATE TABLE IF NOT EXISTS admin_settings (
-                setting_name TEXT PRIMARY KEY,
-                setting_value INTEGER
-            )
-        ''')
-        default_settings = [
-            ('sound_enabled', 1),
-            ('notifications_enabled', 1),
-            ('dark_mode_enabled', 1),
-            ('task_manager_enabled', 1),
-            ('statistics_enabled', 1)
-        ]
-        cursor.executemany('INSERT OR IGNORE INTO admin_settings (setting_name, setting_value) VALUES (?, ?)', default_settings)
-        conn.commit()
-
-def generate_otp():
-    global otp_code, otp_expiration
-    otp_code = ''.join(random.choices(string.digits, k=6))
-    otp_expiration = datetime.now().timestamp() + 300  # OTP valid for 5 minutes
-
-def validate_otp(input_otp):
-    global otp_code, otp_expiration
-    if otp_code and otp_expiration and datetime.now().timestamp() < otp_expiration:
-        return input_otp == otp_code
-    return False
+def get_db():
+    db = sqlite3.connect('database.db')
+    db.row_factory = sqlite3.Row
+    return db
 
 @app.route('/')
 def index():
@@ -92,7 +30,7 @@ def start_session():
     session_type = request.json.get('session_type')
     duration = request.json.get('duration')
     user_id = session.get('user_id')
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('INSERT INTO sessions (user_id, start_time, session_type, duration) VALUES (?, ?, ?, ?)', 
                        (user_id, start_time, session_type, duration))
@@ -104,7 +42,7 @@ def start_session():
 def end_session():
     session_id = request.json.get('session_id')
     end_time = datetime.now().isoformat()
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('UPDATE sessions SET end_time = ? WHERE id = ?', (end_time, session_id))
         conn.commit()
@@ -114,7 +52,7 @@ def end_session():
 def save_script():
     session_id = request.json.get('session_id')
     script_code = request.json.get('script_code')
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('UPDATE sessions SET script_code = ? WHERE id = ?', (script_code, session_id))
         conn.commit()
@@ -123,17 +61,17 @@ def save_script():
 @app.route('/scripts', methods=['GET'])
 def get_scripts():
     user_id = session.get('user_id')
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('SELECT id, start_time, end_time, script_code, rating, comments, session_type, duration FROM sessions WHERE user_id = ?', (user_id,))
         scripts = cursor.fetchall()
-    return jsonify(scripts)
+    return jsonify([dict(row) for row in scripts])
 
 @app.route('/rate_script', methods=['POST'])
 def rate_script():
     session_id = request.json.get('session_id')
     rating = request.json.get('rating')
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('UPDATE sessions SET rating = ? WHERE id = ?', (rating, session_id))
         conn.commit()
@@ -143,7 +81,7 @@ def rate_script():
 def comment_script():
     session_id = request.json.get('session_id')
     comments = request.json.get('comments')
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('UPDATE sessions SET comments = ? WHERE id = ?', (comments, session_id))
         conn.commit()
@@ -158,7 +96,7 @@ def generate_otp_route():
 def reset_database():
     input_otp = request.json.get('otp')
     if validate_otp(input_otp):
-        with sqlite3.connect('database.db') as conn:
+        with get_db() as conn:
             cursor = conn.cursor()
             cursor.execute('DROP TABLE IF EXISTS sessions')
             cursor.execute('DROP TABLE IF EXISTS users')
@@ -177,7 +115,7 @@ def register():
     password = request.json.get('password')
     hashed_password = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt())
     
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         try:
             cursor.execute('INSERT INTO users (username, password) VALUES (?, ?)', (username, hashed_password))
@@ -191,15 +129,15 @@ def login():
     username = request.json.get('username')
     password = request.json.get('password')
     
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('SELECT id, password, is_admin FROM users WHERE username = ?', (username,))
         user = cursor.fetchone()
         
-        if user and bcrypt.checkpw(password.encode('utf-8'), user[1]):
-            session['user_id'] = user[0]
-            session['is_admin'] = user[2]
-            return jsonify({'message': 'Login successful', 'is_admin': user[2]}), 200
+        if user and bcrypt.checkpw(password.encode('utf-8'), user['password']):
+            session['user_id'] = user['id']
+            session['is_admin'] = user['is_admin']
+            return jsonify({'message': 'Login successful', 'is_admin': user['is_admin']}), 200
         else:
             return jsonify({'error': 'Invalid username or password'}), 401
 
@@ -213,7 +151,7 @@ def add_task():
     user_id = session.get('user_id')
     description = request.json.get('description')
     
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('INSERT INTO tasks (user_id, description) VALUES (?, ?)', (user_id, description))
         conn.commit()
@@ -225,19 +163,19 @@ def add_task():
 def get_tasks():
     user_id = session.get('user_id')
     
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('SELECT id, description, completed FROM tasks WHERE user_id = ?', (user_id,))
         tasks = cursor.fetchall()
     
-    return jsonify([{'id': task[0], 'description': task[1], 'completed': bool(task[2])} for task in tasks])
+    return jsonify([dict(row) for row in tasks])
 
 @app.route('/update_task', methods=['POST'])
 def update_task():
     task_id = request.json.get('task_id')
     completed = request.json.get('completed')
     
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('UPDATE tasks SET completed = ? WHERE id = ?', (completed, task_id))
         conn.commit()
@@ -247,7 +185,7 @@ def update_task():
 @app.route('/get_settings', methods=['GET'])
 def get_settings():
     user_id = session.get('user_id')
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('SELECT setting_name, setting_value FROM settings WHERE user_id = ?', (user_id,))
         settings = dict(cursor.fetchall())
@@ -259,7 +197,7 @@ def update_settings():
     user_id = session.get('user_id')
     new_settings = request.json
     
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         for setting, value in new_settings.items():
             cursor.execute('INSERT OR REPLACE INTO settings (user_id, setting_name, setting_value) VALUES (?, ?, ?)', 
@@ -272,7 +210,7 @@ def update_settings():
 def get_statistics():
     user_id = session.get('user_id')
     
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('''
             SELECT 
@@ -285,9 +223,9 @@ def get_statistics():
         stats = cursor.fetchone()
     
     return jsonify({
-        'total_sessions': stats[0],
-        'total_work_time': stats[1],
-        'average_rating': stats[2]
+        'total_sessions': stats['total_sessions'],
+        'total_work_time': stats['total_work_time'],
+        'average_rating': stats['average_rating']
     })
 
 @app.route('/update_admin_settings', methods=['POST'])
@@ -297,7 +235,7 @@ def update_admin_settings():
     
     settings = request.json
     
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         for setting, value in settings.items():
             cursor.execute('INSERT OR REPLACE INTO admin_settings (setting_name, setting_value) VALUES (?, ?)', (setting, int(value)))
@@ -310,7 +248,7 @@ def get_admin_settings():
     if not session.get('is_admin'):
         return jsonify({'error': 'Unauthorized'}), 403
     
-    with sqlite3.connect('database.db') as conn:
+    with get_db() as conn:
         cursor = conn.cursor()
         cursor.execute('SELECT setting_name, setting_value FROM admin_settings')
         settings = dict(cursor.fetchall())