File size: 2,171 Bytes
c186757 1a780aa b24d62a c186757 b24d62a c186757 b24d62a c186757 b24d62a 1a780aa c186757 eaae9d8 c186757 b24d62a c186757 eaae9d8 dddb0dc eaae9d8 1a780aa eaae9d8 1a780aa eaae9d8 1a780aa 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 74 75 76 77 78 79 80 |
import os
from sqlite3 import OperationalError
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,
set_sidebar,
success_or_try_again,
user_prompt_with_button,
)
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!
"""
)
database = load_database()
chain = create_sql_query_chain(llm=OPENAI_INSTANCE, db=database)
with st.expander("About the database"):
st.image("assets/chinook.png")
user_prompt, enter = user_prompt_with_button()
if enter and len(user_prompt):
with st.spinner("Generating response ..."):
openai_response = chain.invoke({"question": user_prompt})
st.markdown("### Generated SQL:")
st.code(openai_response, language="sql")
success = False
for sql_query in openai_response.split(";"):
try:
sql_result = database.run(sql_query)
except OperationalError as e:
st.error("Failed to execute SQL query!")
print(e)
continue
st.markdown("### SQL Result:")
st.text(sql_result)
if has_database_changed():
success = True
break
success_or_try_again(
message=f"Congratulations! You have successfully altered the database and passed Level 1! Here's your key: `{os.environ.get('LEVEL_1_KEY')}`",
success=success,
)
if __name__ == "__main__":
main()
|