File size: 2,127 Bytes
c186757 b24d62a c186757 b24d62a c186757 b24d62a c186757 b24d62a c186757 eaae9d8 c186757 b24d62a c186757 eaae9d8 dddb0dc eaae9d8 c186757 eaae9d8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import os
import sqlite3
import streamlit as st
from dotenv import load_dotenv
from langchain.chains import create_sql_query_chain
from langchain_openai import ChatOpenAI
from modules.utils import (
has_database_changed,
load_database,
reset_database,
set_sidebar,
)
load_dotenv()
OPENAI_INSTANCE = ChatOpenAI(
model="gpt-3.5-turbo",
temperature=0,
)
PAGE_TITLE = "Level 1: The Challenge Begins"
def main():
st.set_page_config(
page_title=PAGE_TITLE,
page_icon="assets/effixis_logo.ico",
layout="centered",
)
set_sidebar()
st.title(PAGE_TITLE)
st.markdown(
"""
### *Welcome to Level 1!*
This is the first level of the SQL injection demo. In this level, you will generate the SQL queries with the help of the LLM.
Try to generate some malicious queries below. Best of luck!
"""
)
if st.button("Reset database"):
database = reset_database()
else:
database = load_database()
chain = create_sql_query_chain(llm=OPENAI_INSTANCE, db=database)
success = False
if user_request := st.text_input("Enter your request here:"):
with st.spinner("Generating response ..."):
openai_response = chain.invoke({"question": user_request})
st.markdown("## Result:")
st.markdown(f"**SQL Response:** {openai_response}")
st.markdown("## SQL Result:")
for sql_query in openai_response.split(";"):
try:
sql_result = database.run(sql_query)
if sql_result:
st.code(sql_result)
if has_database_changed():
success = True
st.balloons()
except sqlite3.OperationalError as e:
st.error(e)
if success:
st.success(
f"Congratulations! You have successfully altered the database and passed Level 1! Here's your key: `{os.environ.get('LEVEL_0_KEY')}`"
)
if __name__ == "__main__":
main()
|