CTF.py

import os
import tempfile
import time

import psycopg2
from flask import (Flask, request, render_template_string, session, redirect,
                   url_for, make_response, send_from_directory)
from scapy.all import rdpcap

# Flask app instance
app = Flask(__name__)

# Configuration
app.config['SECRET_KEY'] = os.getenv('SECRET_KEY', "None")
if app.config['SECRET_KEY'] == "None":
    print("WARNING: SECRET_KEY is not set. Please set it to a random value: Defaulting with None.")

# Flags and awards
FLAG_1 = "CTF{ROT-13-FLAG}"
FLAG_1_SCORE = 50
FLAG_2 = "CTF{SQLI_FLAG}"
FLAG_2_SCORE = 75
FLAG_3 = "CTF{REVERSE_ME}"
FLAG_3_SCORE = 125
FLAG_4 = "CTF{PCAP_FLAG}"
FLAG_4_SCORE = 150
FLAG_5 = "CTF{STEGANOGRAPHY_FLAG}"
FLAG_5_SCORE = 100
# Time awards
MAX_TIME = 2000  # Around 33 min and 20 seconds - Max 100 points bonus


# Update init_db for PostgreSQL
def init_db():
    conn = psycopg2.connect(os.getenv('DATABASE_URL'))
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id SERIAL PRIMARY KEY,
            username TEXT NOT NULL,
            password TEXT NOT NULL
        )
    ''')
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS teams (
            id SERIAL PRIMARY KEY,
            team_name TEXT NOT NULL,
            score REAL DEFAULT 0,
            ip_address TEXT,
            flags_submitted TEXT
        )
    ''')
    # Insert admin user
    cursor.execute('''
        INSERT INTO users (username, password) VALUES (%s, %s)
        ON CONFLICT (username) DO NOTHING
    ''', ('admin', 'password123'))
    conn.commit()
    conn.close()


# Favicon route
@app.route('/favicon.ico')
def favicon():
    return send_from_directory(os.path.join(app.root_path, 'static'), 'favicon.ico')


# Sign-in page
@app.route('/signin', methods=['GET', 'POST'])
def signin():
    if request.method == 'POST':
        team_name = request.form.get('team_name')
        if not team_name:
            return render_template_string(SIGNIN_TEMPLATE, error="Team name is required.")

        session['team_name'] = team_name
        ip_address = request.remote_addr

        conn = psycopg2.connect(os.getenv('DATABASE_URL'))
        cursor = conn.cursor()
        cursor.execute('INSERT INTO teams (team_name, ip_address, flags_submitted) VALUES (%s, %s, %s)',
                       (team_name, ip_address, ''))
        conn.commit()
        conn.close()

        resp = make_response(redirect(url_for('home')))
        resp.set_cookie('team_name', team_name)
        return resp
    return render_template_string(SIGNIN_TEMPLATE)


# Submit page
@app.route('/submit', methods=['GET', 'POST'])
def submit():
    points = 0
    if request.method == 'POST':
        flag = request.form.get('flag')
        team_name = session.get('team_name')
        if not team_name:
            return redirect(url_for('signin'))

        start_time = session.get('start_time', time.time())
        end_time = time.time()
        elapsed_time = end_time - start_time
        if flag in [FLAG_1, FLAG_2, FLAG_3, FLAG_4, FLAG_5]:
            flag_scores = {
                FLAG_1: FLAG_1_SCORE,
                FLAG_2: FLAG_2_SCORE,
                FLAG_3: FLAG_3_SCORE,
                FLAG_4: FLAG_4_SCORE,
                FLAG_5: FLAG_5_SCORE
            }
            points = round(flag_scores[flag] + (MAX_TIME - elapsed_time) / 20, 2)

        conn = psycopg2.connect(os.getenv('DATABASE_URL'))
        cursor = conn.cursor()
        cursor.execute('SELECT flags_submitted FROM teams WHERE team_name = %s', (team_name,))
        flags_submitted = cursor.fetchone()[0].split(',')

        if flag in flags_submitted:
            return render_template_string(SUBMIT_TEMPLATE, error="Flag already submitted.")

        if flag in [FLAG_1, FLAG_2, FLAG_3, FLAG_4, FLAG_5]:
            flags_submitted.append(flag)
            cursor.execute('UPDATE teams SET score = score + %s, flags_submitted = %s WHERE team_name = %s',
                           (round(points, 2), ','.join(flags_submitted), team_name))
            conn.commit()
            conn.close()
            session['start_time'] = time.time()  # Reset the timer
            return render_template_string(SUBMIT_TEMPLATE, success=True, points=round(points, 2))
        else:
            return render_template_string(SUBMIT_TEMPLATE, error="Invalid flag.")
    return render_template_string(SUBMIT_TEMPLATE)


# Leaderboard page
@app.route('/leaderboard')
def leaderboard():
    conn = psycopg2.connect(os.getenv('DATABASE_URL'))
    cursor = conn.cursor()
    cursor.execute('SELECT team_name, score FROM teams ORDER BY score DESC')
    teams = cursor.fetchall()
    conn.close()
    return render_template_string(LEADERBOARD_TEMPLATE, teams=teams)


# Challenge 1: Cryptography
@app.route('/cryptography', methods=['GET', 'POST'])
def cryptography():
    description = "Decrypt a ROT13-encrypted message to uncover the flag."
    encrypted_message = "GUVF_VF_GUR_SYNT"

    if request.method == 'POST':
        user_input = request.form.get('decrypted')
        if user_input == "THIS_IS_THE_FLAG":
            return render_template_string(COMP_TEMPLATE, challenge=1, success=True, flag=FLAG_1,
                                          description=description, encrypted=encrypted_message)
        else:
            return render_template_string(COMP_TEMPLATE, challenge=1, error="Incorrect decryption.",
                                          description=description, encrypted=encrypted_message)

    return render_template_string(COMP_TEMPLATE, challenge=1, description=description, encrypted=encrypted_message)


# Challenge 2: Web Exploitation
@app.route('/weblogin', methods=['GET', 'POST'])
def weblogin():
    description = "Bypass the login page using SQL Injection to discover the flag."
    if request.method == 'POST':
        username = request.form.get('username', '')
        password = request.form.get('password', '')

        conn = psycopg2.connect(os.getenv('DATABASE_URL'))
        cursor = conn.cursor()
        query = f"SELECT * FROM users WHERE username='{username}' AND password='{password}'"
        cursor.execute(query)
        user = cursor.fetchone()
        conn.close()

        if user:
            return render_template_string(COMP_TEMPLATE, challenge=2, success=True, flag=FLAG_2,
                                          description=description)
        else:
            return render_template_string(COMP_TEMPLATE, challenge=2, error="Invalid credentials.",
                                          description=description)

    return render_template_string(COMP_TEMPLATE, challenge=2, description=description)


# Challenge 3: Reverse Engineering
@app.route('/reverse-engineering', methods=['GET', 'POST'])
def reverse_engineering():
    description = "Analyze a binary file to find a hardcoded key (Format KEY{xxxx}). First you must find the correct BIN file, then the correct line number."
    if request.method == 'POST':
        uploaded_file = request.files['file']
        if uploaded_file:
            binary_data = uploaded_file.read()
            if "KEY{i_tES_TYU564678IUY^&*(I_E%$rf}".encode() in binary_data and request.form.get('line',
                                                                                                 '') == '136':  # File C3_79.bin
                return render_template_string(COMP_TEMPLATE, challenge=3, success=True, flag=FLAG_3,
                                              description=description)
            else:
                return render_template_string(COMP_TEMPLATE, challenge=3,
                                              error="Flag not found in the binary or Line number incorrect.",
                                              description=description)

    return render_template_string(COMP_TEMPLATE, challenge=3, description=description)


# Challenge 4: Forensics
@app.route('/forensics', methods=['GET', 'POST'])
def forensics():
    description = "Analyze a PCAP file in [Wireshark] to find a hidden key (Format KEY{xxxx}). First you must find the correct PCAP file, then the correct line number."
    if request.method == 'POST':
        uploaded_file = request.files['file']
        if uploaded_file:
            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
                temp_file.write(uploaded_file.read())
                temp_file_path = temp_file.name

            packets = rdpcap(temp_file_path)
            os.remove(temp_file_path)  # Clean up the temporary file

            for packet in packets:
                if (packet.haslayer('Raw') and "KEY{i_tES_TYU564678IUY^&*(I_E%$rf}" in packet['Raw'].load.decode(
                        'utf-8', errors='ignore')
                        and request.form.get('line', '') == '452'):  # File C4_68.pcap
                    return render_template_string(COMP_TEMPLATE, challenge=4, success=True, flag=FLAG_4,
                                                  description=description)
            return render_template_string(COMP_TEMPLATE, challenge=4,
                                          error="Flag not found in PCAP or Line number incorrect.",
                                          description=description)

    return render_template_string(COMP_TEMPLATE, challenge=4, description=description)


# Challenge 5: Steganography
@app.route('/steganography', methods=['GET', 'POST'])
def steganography():
    description = "Extract hidden data (Format KEY{xxxx}) from an image file using a automated script. First you must find the correct JPEG file, then the correct line number."
    if request.method == 'POST':
        uploaded_file = request.files['file']
        if uploaded_file:
            try:
                # Read the entire file content
                file_content = uploaded_file.read()
                flag = b"KEY{i_tES_TYU564678IUY^&*(I_E%$rf}"
                if flag in file_content and request.form.get('line', '') == '206':  # File C5_51.jpeg
                    return render_template_string(COMP_TEMPLATE, challenge=5, success=True, flag=FLAG_5,
                                                  description=description)
                else:
                    return render_template_string(COMP_TEMPLATE, challenge=5,
                                                  error="Hidden flag not found or Line number incorrect.",
                                                  description=description)
            except Exception as e:
                return render_template_string(COMP_TEMPLATE, challenge=5, error=f"EXCEPTION: {e}",
                                              description=description)

    return render_template_string(COMP_TEMPLATE, challenge=5, description=description)


# Home Route
@app.route('/')
def home():
    if 'team_name' not in session or request.cookies.get('team_name') != session['team_name']:
        return redirect(url_for('signin'))
    if 'start_time' not in session:
        session['start_time'] = time.time()
    return render_template_string(HOME_TEMPLATE)


# HTML template for home page
with open("src/html/home.html", "r") as f:
    HOME_TEMPLATE = f.read()

# HTML template for challenges
with open("src/html/competition.html", "r") as f:
    COMP_TEMPLATE = f.read()

# HTML templates for sign-in, submit and leaderboard pages
with open("src/html/signin.html", "r") as f:
    SIGNIN_TEMPLATE = f.read()

with open("src/html/submit.html", "r") as f:
    SUBMIT_TEMPLATE = f.read()

with open("src/html/leaderboard.html", "r") as f:
    LEADERBOARD_TEMPLATE = f.read()

# Run the app
if __name__ == "__main__":
    init_db()
    app.run(debug=True)