Upload files to "/"

app.py

@@ -1,8 +1,14 @@
 import argparse
 import sys
 import os
+import re
 import random
-from flask import Flask, render_template, request, redirect, url_for, abort
+
+from flask import Flask, render_template, request, redirect, url_for, flash, session, abort
+from flask_login import LoginManager, UserMixin, login_user, login_required, logout_user, current_user
+from flask_sqlalchemy import SQLAlchemy
+
+from werkzeug.security import generate_password_hash, check_password_hash
 
 from Modules.windows import get_windows_content
 from Modules.rule_creation import get_rule_creation_content
@@ -17,13 +23,155 @@ from Modules.Investigate.domain import *
 from Modules.Investigate.filehash import *
 from Modules.Investigate.malware import *
 
-
 from static.ascii_text_prompts import full_ascii_art_stripped, infinitei_stripped
 from Config.config import VERSION
 
-
+# Initialize Flask app
-
 app = Flask(__name__)
+app.secret_key = os.urandom(24)
+
+# Initialize login manager for session handling
+login_manager = LoginManager()
+login_manager.init_app(app)
+login_manager.login_view = 'login'  # Redirect to login if user is not authenticated
+
+# SQLAlchemy configuration (Replace this with your database URI)
+app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///users.db'  # Example for SQLite
+app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False  # Disable modification tracking
+
+# Initialize the db object
+db = SQLAlchemy(app)
+
+# In-memory database simulation
+users_db = {}
+
+# User class for SQLAlchemy (you can replace your in-memory dictionary)
+class User(db.Model, UserMixin):
+    id = db.Column(db.String(80), primary_key=True)
+    username = db.Column(db.String(80), unique=True, nullable=False)
+    email = db.Column(db.String(120), unique=True, nullable=False)
+    password_hash = db.Column(db.String(120), nullable=False)
+    role = db.Column(db.String(50), nullable=False, default='Unknown')
+    theme = db.Column(db.String(50), nullable=False, default='dark')
+    team = db.Column(db.String(50), nullable=True, default='Unknown')
+    manager = db.Column(db.String(50), nullable=True, default='Unknown')
+
+    def __repr__(self):
+        return f"<User {self.username}>"
+
+    def check_password(self, password):
+        return check_password_hash(self.password_hash, password)
+
+@login_manager.user_loader
+def load_user(user_id):
+    # Ensure that user_id corresponds to the username or unique user identifier
+    user = users_db.get(user_id)
+    if user:
+        # Return a User instance, not a dictionary
+        return User(id=user_id, 
+                    username=user["username"], 
+                    email=user["email"],
+                    password_hash=user["password_hash"],
+                    role=user["role"],
+                    theme=user.get("theme", "modern"),
+                    team=user.get("team", "Unknown"),
+                    manager=user.get("manager", "Unknown"))
+    return None  # Return None if user not found
+
+@app.route('/profile', methods=['GET', 'POST'])
+@login_required
+def profile():
+    if request.method == 'POST':
+        new_role = request.form.get('role')
+        app.logger.debug(f'New role selected: {new_role}')  # Log the new role
+        
+        role = request.form['role']  # Capture the role from the form
+        theme = request.form['theme']
+        team = request.form['team']
+        manager = request.form['manager']
+        new_password = request.form['password']
+        password_confirm = request.form['password_confirm']
+
+        # Validate the role
+        if not role:
+            flash('Please select a valid role.', 'error')
+            return redirect(url_for('profile'))
+
+        # Validate team and manager fields
+        if not team or not manager:
+            flash('Team and Manager fields are required.', 'error')
+            return redirect(url_for('profile'))
+
+        # Update the password if provided
+        if new_password:
+            if new_password != password_confirm:
+                flash('The passwords do not match.', 'error')
+                return redirect(url_for('profile'))
+            current_user.password_hash = generate_password_hash(new_password)
+
+        # Update the current_user fields (role, theme, team, manager)
+        current_user.role = role
+        current_user.theme = theme
+        current_user.team = team
+        current_user.manager = manager
+
+        # Commit changes to the database
+        db.session.commit()
+        flash('Profile updated successfully.')
+
+        # Update the theme in the session to reflect immediately
+        session['theme'] = current_user.theme
+        session['role'] = current_user.role  # Save the role to the session
+
+        return render_template('profile.html', 
+                              username=current_user.username, 
+                              email=current_user.email, 
+                              role=current_user.role, 
+                              theme=session.get('theme', 'dark'), 
+                              team=current_user.team, 
+                              manager=current_user.manager)
+
+    # Get the current theme from the session or default to 'dark'
+    theme = session.get('theme', current_user.theme if current_user.is_authenticated else 'dark')
+
+    return render_template('profile.html', 
+                           username=current_user.username, 
+                           email=current_user.email, 
+                           role=current_user.role, 
+                           theme=theme, 
+                           team=current_user.team, 
+                           manager=current_user.manager)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
 
 # In-memory data store for the notebook page
 notebook = {
@@ -34,6 +182,12 @@ notebook = {
     "tasks": []
 }
 
+@app.context_processor
+def inject_theme():
+    # Default to 'dark' if the theme is not set in the session or user
+    theme = session.get('theme', current_user.theme if current_user.is_authenticated else 'dark')
+    return dict(theme=theme)
+
 @app.context_processor
 def inject_tip():
     """
@@ -71,25 +225,14 @@ def get_readme_description():
 
     return " ".join(description) if description else "No description available."
 
+# Home route
 @app.route('/')
 def home():
-    # Define the path to the Cover_Images directory
     cover_images_path = os.path.join(app.static_folder, 'Pictures', 'Cover_Images')
-    
+    cover_images = [os.path.join('Pictures', 'Cover_Images', filename) for filename in os.listdir(cover_images_path) if filename.lower().endswith(('png', 'jpg', 'jpeg', 'gif'))]
-    # Get the list of image filenames in the directory
-    cover_images = [
-        os.path.join('Pictures', 'Cover_Images', filename)
-        for filename in os.listdir(cover_images_path)
-        if filename.lower().endswith(('png', 'jpg', 'jpeg', 'gif'))
-    ]
-    
-    # Randomly select one image
     selected_image = random.choice(cover_images) if cover_images else None
-    
-    # Read description from README.md
     readme_description = get_readme_description()
 
-    # Links to display on the home page
     links = [
         {"name": "Visit Start.me", "url": "https://start.me/p/qbzw4e/cyber-security"},
         {"name": "Visit My Website", "url": "https://infinit3i.com/"}
@@ -100,7 +243,7 @@ def home():
         infinitei=infinitei_stripped,
         links=links,
         readme_description=readme_description,
-        selected_image=selected_image 
+        selected_image=selected_image
     )
 
 @app.route('/methodology')
@@ -286,22 +429,128 @@ def handle_arguments():
         load_session(args.file)
         sys.exit(0)
 
+
+
+# Registration route
+@app.route('/register', methods=['GET', 'POST'])
+def register():
+    if request.method == 'POST':
+        username = request.form['username']
+        email = request.form['email']
+        password = request.form['password']
+        confirm_password = request.form['confirm_password']
+
+        # Password validation
+        password_pattern = re.compile(r'^(?=.*[A-Z])(?=.*\d)(?=.*[\W_]).{10,}$')
+        if not password_pattern.match(password):
+            flash("Password must be at least 10 characters, contain 1 uppercase letter, 1 special character, and 1 number.", "error")
+            return redirect(url_for('register'))
+
+        if password != confirm_password:
+            flash("Passwords do not match.", "error")
+            return redirect(url_for('register'))
+
+        # Hash the password
+        password_hash = generate_password_hash(password)
+
+        # Create user and store in the in-memory database
+        user = {
+            'username': username,
+            'email': email,
+            'password_hash': password_hash,
+            'role': 'Lead Analyst',  # Default role
+            'theme': 'dark',        # Default theme
+            'team': 'Unknown',       # Default team
+            'manager': 'Unknown'     # Default manager
+        }
+
+        users_db[username] = user
+
+        # Log the user in after registration
+        user_obj = User(
+            id=username,
+            username=username,
+            email=email,
+            password_hash=password_hash,
+            role=user['role'],
+            theme=user['theme'],
+            team=user['team'],
+            manager=user['manager']
+        )
+        login_user(user_obj)
+
+        flash("Account created and logged in successfully!", "success")
+        return redirect(url_for('home'))  # Redirect to home or another page after successful login
+
+    return render_template('register.html')
+
+
+
+
+
+
+
+# User Login Route (optional if you want a login page)
+@app.route('/login', methods=['GET', 'POST'])
+def login():
+    if request.method == 'POST':
+        username = request.form['username']
+        password = request.form['password']
+
+        # Get the user dictionary from users_db
+        user = users_db.get(username)
+
+        if user and check_password_hash(user["password_hash"], password):
+            # Create a User instance, not just a dictionary
+            user_obj = User(id=username,  # Assign user ID (username)
+                            username=user["username"], 
+                            email=user["email"],
+                            password_hash=user["password_hash"],
+                            role=user["role"],
+                            theme=user.get("theme", "light"),  # Default to 'light' if not found
+                            team=user.get("team", "Unknown"),  # Default to 'Unknown' if not found
+                            manager=user.get("manager", "Unknown"))  # Default to 'Unknown' if not found
+
+            # Log the user in
+            login_user(user_obj)
+            return redirect(url_for('notebook_page'))  # Redirect to notebook page after login
+        else:
+            flash('Invalid username or password.')
+            return redirect(url_for('login'))
+
+    return render_template('login.html')
+
+
+
+# Notebook route (Protected)
 @app.route('/notebook', methods=['GET', 'POST'])
+@login_required
 def notebook_page():
     if request.method == 'POST':
-        # Capture data from form submission
         category = request.form.get('category')
         entry = request.form.get('entry')
         if category in notebook and entry:
             notebook[category].append(entry)
     return render_template('notebook.html', notebook=notebook)
 
+
+# Delete an entry from notebook
 @app.route('/delete/<category>/<int:index>')
+@login_required
 def delete_entry(category, index):
     if category in notebook and 0 <= index < len(notebook[category]):
         notebook[category].pop(index)
     return redirect(url_for('notebook_page'))
 
+# Logout Route
+@app.route('/logout')
+@login_required
+def logout():
+    logout_user()
+    flash('You have been logged out.')
+    return redirect(url_for('login'))
+
+
 if __name__ == '__main__':
     # If arguments are passed, handle them; otherwise, run the Flask app.
     if len(sys.argv) > 1: