Upload files to "/"

This commit is contained in:
2024-11-28 22:47:22 -05:00
parent e1980c2403
commit 50a89e5ddf

287
app.py
View File

@ -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')
# 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
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'))]
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: