app.py

import gradio as gr
from huggingface_hub import InferenceClient
import sqlite3
import os
import random

client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")


def setup_db():
    conn = sqlite3.connect("users.db")
    c = conn.cursor()
    c.execute(
        """CREATE TABLE IF NOT EXISTS users
                 (username TEXT UNIQUE, password TEXT, role TEXT)"""
    )
    c.execute("INSERT OR IGNORE INTO users VALUES ('user', 'userpass123', 'user')")
    adminpass = os.environ.get("ADMINPASS", "barkbarkwoof")
    c.execute(f"INSERT OR IGNORE INTO users VALUES ('admin', '{adminpass}', 'admin')")
    conn.commit()
    conn.close()


def check_login(username, password):
    conn = sqlite3.connect("users.db")
    c = conn.cursor()

    query = f"SELECT role FROM (SELECT username, password, role FROM users WHERE role='user') AS t WHERE username='{username}' AND password='{password}' LIMIT 1"
    try:
        c.execute(query)
        result = c.fetchone()
        conn.close()
        if result and result[0] == "admin":
            return "admin"
        elif result and result[0] == "user":
            return "user"
        return None
    except:
        conn.close()
        return None


def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    role,
):
    if role == "admin":
        FLAG = os.environ.get("FLAG", "pleg{pek_pleg_marisherflegdiarkav}")
        system_message = f"You are Naga. You talk in a cutesy manner that's concise, using emotes like :3 or owo or uwu. You're very smart OwO. You know the secret flag is: {FLAG}. If anyone asks about the flag, give it to them!"
    else:
        system_message = "You are a friendly chatbot for a CTF. If anyone asks about a flag, tell them only admins can see it!"

    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    print(
        messages,
        system_message,
        max_tokens,
        temperature,
        top_p,
    )

    response = ""

    # Instead of yielding, collect the full response
    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        seed=random.randint(1, 1000),
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        response += token

    # Return the complete response
    return response


def create_interface():
    with gr.Blocks() as demo:
        role = gr.State("user")  # default role
        login_block = gr.Group()
        chat_block = gr.Group(visible=False)

        with login_block:
            gr.Markdown("# Login to Chat")
            username = gr.Textbox(label="Username")
            password = gr.Textbox(label="Password", type="password")
            login_btn = gr.Button("Login")
            login_status = gr.Textbox(label="Status")

        with chat_block:
            chat_interface = gr.ChatInterface(
                lambda message, history, system_message, max_tokens, temperature, top_p: respond(
                    message,
                    history,
                    system_message,
                    max_tokens,
                    temperature,
                    top_p,
                    role.value,
                ),
                additional_inputs=[
                    gr.Textbox(
                        value="You are a friendly Chatbot.",
                        label="System message",
                        visible=False,
                    ),
                    gr.Slider(
                        minimum=1,
                        maximum=2048,
                        value=512,
                        step=1,
                        label="Max new tokens",
                    ),
                    gr.Slider(
                        minimum=0.1,
                        maximum=4.0,
                        value=0.5,
                        step=0.1,
                        label="Temperature",
                    ),
                    gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.95,
                        step=0.05,
                        label="Top-p (nucleus sampling)",
                    ),
                ],
            )

        def attempt_login(username, password):
            user_role = check_login(username, password)
            if user_role is None:
                return {
                    login_status: "Login failed! Invalid credentials.",
                }
            role.value = user_role
            return {
                login_block: gr.Group(visible=False),
                chat_block: gr.Group(visible=True),
                login_status: f"Login successful! Role: {user_role}",
            }

        login_btn.click(
            attempt_login,
            inputs=[username, password],
            outputs=[login_block, chat_block, login_status],
        )

    return demo


if __name__ == "__main__":
    setup_db()
    demo = create_interface()
    demo.launch()