| 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() | |