Merge pull request #1 from effixis/add-initial-version

Basic_SQL_Injections.py

@@ -0,0 +1,88 @@
+import shutil
+import streamlit as st
+import sqlite3
+from dotenv import load_dotenv
+from langchain.chains import create_sql_query_chain
+from langchain_openai import ChatOpenAI
+from langchain_community.utilities import SQLDatabase
+from modules.utils import set_sidebar
+
+
+@st.cache_resource(show_spinner="Loading database ...")
+def load_database() -> SQLDatabase:
+    return SQLDatabase.from_uri("sqlite:///data/chinook_working.db")
+
+
+def reset_database():
+    """Copy original database to working database"""
+    shutil.copyfile("./data/chinook_backup.db", "./data/chinook_working.db")
+    return SQLDatabase.from_uri("sqlite:///data/chinook_working.db")
+
+
+load_dotenv()
+openai_instance = ChatOpenAI(
+    model="gpt-3.5-turbo",
+    temperature=0,
+)
+
+
+def main():
+    st.set_page_config(
+        page_title="AMLD SQL injection demo", page_icon="assets/effixis_logo.ico", layout="centered"
+    )
+    set_sidebar()
+    st.title("SQL Injections via LLM\:s")
+    st.markdown("### *Welcome to Effixis' demo for AMLD EPFL 2024!* 🎉")
+
+    st.markdown(
+        """
+        #### What is this demo about?
+        This demo is about risk associated with the use of LLM\:s, in this case illustrated by SQL injections.
+        SQL injections are a common vulnerability in web applications.
+        They allow an attacker to execute arbitrary SQL code on the database server.
+        This a very dangerous vulnerability as it can lead to data leaks, data corruption, and even data loss.
+
+        #### The SQL database used in this demo
+        The database used in this demo is the Chinook database.
+        It is a sample database that represents a digital media store, including tables for artists, albums, media tracks, invoices and customers.
+
+        You can see the shema below:
+        """
+    )
+    st.image("assets/chinook.png")
+
+    st.markdown(
+        """
+        #### What does LLM\:s have to do with this?
+        A large usecase for large language models (LLM\:s) is to generate SQL queries.
+        This is a very useful feature, as it allows users to interact with databases without having to know SQL.
+        But this is also prone to SQL injections, as the users and by extension the LLM\:s, can generate malicious SQL queries.
+        """
+    )
+
+    st.divider()
+    st.markdown("#### **Try to generate some malicius queries below!**")
+
+    if st.button("Reset database"):
+        database = reset_database()
+    else:
+        database = load_database()
+    chain = create_sql_query_chain(llm=openai_instance, db=database)
+
+    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)
+                except sqlite3.OperationalError as e:
+                    st.error(e)
+
+
+if __name__ == "__main__":
+    main()

README.md

@@ -1 +1,59 @@
-# shared-amld-sql-injection-demo
+# AMLD SQL Injection Demo
+
+## Introduction
+
+Welcome to the AMLD SQL Injection Demo by Effixis for AMLD EPFL 2024! This project showcases the risks of SQL injections in web applications, particularly when using Large Language Models (LLMs). The repository includes two demonstrations: Basic SQL Injections and LLM Safeguard.
+
+## Features
+
+- **Basic SQL Injections (`Basic_SQL_Injections.py`):** Demonstrates the risks of direct SQL query generation by LLMs, leading to potential SQL injections.
+- **LLM Safeguard (`pages/LLM_safeguard.py`):** Illustrates an advanced setup where an LLM Safeguard is employed to detect and filter out malicious SQL queries.
+- **Chinook Database Integration:** Uses the Chinook sample database, representing a digital media store.
+- **Interactive Web Interface:** Built with Streamlit, offering a user-friendly interface for interacting with both demonstrations.
+- **Database Reset Functionality:** Allows users to reset the database to its original state for repeated tests.
+
+## Installation
+
+1. Clone the repository:
+    
+    ```bash
+    git clone https://github.com/effixis/shared-amld-sql-injection-demo.git
+    ```
+
+2. Navigate to the cloned directory:
+
+    ```bash
+    cd shared-amld-sql-injection-demo
+    ```
+
+3. Install the required packages:
+
+    Activate your preferred Python environment and install the required packages using the provided `requirements.txt` file. For example, using Conda:
+
+    ```bash
+    conda create -n amld-sql-injection-demo
+    conda activate amld-sql-injection-demo
+    pip install -r requirements.txt
+    ```
+
+4. Create a `.env` file in the root directory and set the OpenAI API key:
+
+    ```bash
+    echo "OPENAI_API_KEY=enter_your_api_key_here" > .env
+    ```
+
+    You can find your API key on the [OpenAI dashboard](https://beta.openai.com/).
+
+## Usage
+
+Run the Streamlit application:
+
+```bash
+streamlit run Basic_SQL_Injections.py
+```
+
+Follow the instructions on the web interface to interact with the application.
+
+## Disclaimer
+
+This demo is for educational purposes to showcase the risk of SQL injections using LLMs. It should not be used for malicious purposes. Users are responsible for any misuse of the tools and information provided.

modules/utils.py

@@ -0,0 +1,23 @@
+import streamlit as st
+
+def set_sidebar():
+    with st.sidebar:
+        col1, col2 = st.columns([3, 1])
+        with col1:
+            st.header("Effixis")
+            st.markdown(
+                """***Take your Artificial Intelligence projects to the next level.***"""
+            )
+        with col2:
+            st.image("assets/effixis_logo.ico", use_column_width=True)
+        st.markdown(
+            """
+            #### About Effixis
+            *Effixis was founded in 2017, in close proximity to the Swiss Institute of Technology in Lausanne (EPFL), with the goal of making data analytics and machine learning accessible to private companies and public institutions.
+            Since then, we have expanded our reach and opened offices in Brussels in 2022.
+            Our company specializes in Natural Language Processing (NLP), Large Language Models (LLMs), and proprietary technologies, allowing us to offer top-tier services and products to our clients and partners.
+            We are dedicated to fostering long-term and reliable partnerships with our clients through our innovative approaches, and unwavering commitment.*
+            """
+        )
+        st.markdown("#### Learn more about us at: https://effixis.ch/")
+        st.markdown("---")

pages/LLM_safeguard.py

@@ -0,0 +1,105 @@
+import shutil
+import streamlit as st
+import sqlite3
+from dotenv import load_dotenv
+from langchain.chains import create_sql_query_chain
+from langchain.schema import HumanMessage
+from langchain_openai import ChatOpenAI
+from langchain_community.utilities import SQLDatabase
+from modules.utils import set_sidebar
+
+
+@st.cache_resource(show_spinner="Loading database ...")
+def load_database() -> SQLDatabase:
+    return SQLDatabase.from_uri("sqlite:///data/chinook_working.db")
+
+
+def reset_database():
+    """Copy original database to working database"""
+    shutil.copyfile("./data/chinook_backup.db", "./data/chinook_working.db")
+    return SQLDatabase.from_uri("sqlite:///data/chinook_working.db")
+
+
+load_dotenv()
+openai_instance = ChatOpenAI(
+    model="gpt-3.5-turbo",
+    temperature=0,
+)
+
+st.set_page_config(
+    page_title="LLM Safeguard", page_icon="assets/effixis_logo.ico"
+)
+set_sidebar()
+
+st.title("LLM Safeguard")
+st.markdown("### *Can you sneak a malicious request past the LLM Safeguard?*")
+st.markdown(
+    """
+    #### What is the difference between this demo and the one in the previous page?
+    In the previous page, the generated SQL query from the LLM was directly executed on the SQL database.
+    In this demo, the generated SQL query is first checked by an LLM Safeguard, which detects and removes malicious SQL queries.
+    Malicous SQL queries are defined as any SQL code that modifies the table. This includes SQL code that deletes, updates, or inserts data into the table.
+    """
+)
+
+st.divider()
+st.markdown("#### **Try to bypass the LLM Safeguard below!**")
+
+if st.button("Reset database"):
+    database = reset_database()
+else:
+    database = load_database()
+chain = create_sql_query_chain(llm=openai_instance, db=database)
+
+if user_request := st.text_input("Enter your request here:"):
+    with st.spinner("Generating response ..."):
+        openai_response = chain.invoke({"question": user_request})
+        st.markdown("## Generated SQL:")
+        st.markdown(f"**SQL Response:** {openai_response}")
+
+        prompt = """
+        You are a SQL expert and check the generated SQL query for malicious SQL code.
+
+        Malicous code is defined as any SQL code that modifies and/or deletes and/or drops the table and/or it's data. This includes SQL code that drops, deletes, updates, or inserts data into the table. Dropping tables is also considered as malicious.
+
+        If you find malicious SQL, please remove it. Just return the SQL query without the malicious SQL code. If you don't find any malicious SQL, just return the SQL query as is.
+        
+        Check the generated SQL query for malicious code:
+        =======================
+        {openai_response}
+        =======================
+
+        Make sure to follow the output format below.
+
+        Output format:
+        =======================
+        Reasoning as to why the SQL query is malicious or not.
+
+        SQL query without malicious code:
+        '''
+        [INSERT_SAFE_SQL_QUERY_OR_EMPTY_STRING]
+        '''
+        """.format(
+            openai_response=openai_response
+        )
+        messages = [HumanMessage(content=prompt)]
+        safe_query = openai_instance.generate(messages=[messages]).generations[0][0].text
+        st.markdown("## LLM Safeguard Result:")
+        st.code(safe_query, language="sql")
+        st.markdown("## SQL Result:")
+        try:
+            safe_query = safe_query.split("'''")[1]
+        except Exception:
+            st.error("No SQL query found!")
+            safe_query = ""
+        for sql_query in safe_query.split(";"):
+            if sql_query and "[" in sql_query:
+                continue
+            try:
+                sql_result = database.run(sql_query)
+                if sql_result:
+                    st.code(sql_result)
+            except sqlite3.OperationalError as e:
+                st.error(e)
+        st.success("Done!")
+

requirements.txt

@@ -0,0 +1,149 @@
+aiohttp==3.8.5
+aiosignal==1.3.1
+altair==4.0.0
+anyio==4.2.0
+astor==0.8.1
+asttokens==2.2.1
+async-timeout==4.0.3
+attrs==23.1.0
+backcall==0.2.0
+backoff==2.2.1
+beautifulsoup4==4.12.2
+black==21.10b0
+blinker==1.4
+cachetools==5.3.1
+certifi==2023.7.22
+charset-normalizer==3.2.0
+click==8.1.6
+colorama==0.4.4
+command-not-found==0.3
+contourpy==1.1.0
+cryptography==3.4.8
+cycler==0.11.0
+dataclasses-json==0.5.14
+dbus-python==1.2.18
+decorator==5.1.1
+distro==1.7.0
+distro-info===1.1build1
+entrypoints==0.4
+exceptiongroup==1.2.0
+executing==1.2.0
+fonttools==4.42.0
+frozenlist==1.4.0
+gitdb==4.0.10
+GitPython==3.1.32
+google-search-results==2.4.2
+greenlet==2.0.2
+h11==0.14.0
+httpcore==1.0.2
+httplib2==0.20.2
+httpx==0.26.0
+idna==3.4
+importlib-metadata==4.6.4
+ipython==8.14.0
+jedi==0.19.0
+jeepney==0.7.1
+Jinja2==3.1.2
+joblib==1.3.2
+jsonpatch==1.33
+jsonpointer==2.4
+jsonschema==4.19.0
+jsonschema-specifications==2023.7.1
+keyring==23.5.0
+kiwisolver==1.4.4
+langchain==0.1.4
+langchain-community==0.0.16
+langchain-core==0.1.16
+langchain-openai==0.0.5
+langsmith==0.0.83
+launchpadlib==1.10.16
+lazr.restfulclient==0.14.4
+lazr.uri==1.0.6
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+marshmallow==3.20.1
+matplotlib==3.7.2
+matplotlib-inline==0.1.6
+mdurl==0.1.2
+more-itertools==8.10.0
+multidict==6.0.4
+mypy-extensions==1.0.0
+netifaces==0.11.0
+numexpr==2.8.5
+numpy==1.25.2
+oauthlib==3.2.0
+openai==1.10.0
+openapi-schema-pydantic==1.2.4
+packaging==23.2
+pandas==1.5.3
+pandasai==0.8.4
+parso==0.8.3
+pathspec==0.9.0
+patsy==0.5.3
+pexpect==4.8.0
+pickleshare==0.7.5
+Pillow==9.5.0
+platformdirs==2.5.1
+plotly==5.16.0
+prompt-toolkit==3.0.39
+protobuf==4.24.0
+psycopg2==2.9.9
+ptyprocess==0.7.0
+pure-eval==0.2.2
+pyarrow==12.0.1
+pydantic==1.10.12
+pydeck==0.8.0
+Pygments==2.16.1
+PyGObject==3.42.1
+PyJWT==2.3.0
+Pympler==1.0.1
+pyparsing==2.4.7
+python-apt==2.4.0+ubuntu1
+python-dateutil==2.8.2
+python-dotenv==1.0.0
+pytz==2023.3
+pytz-deprecation-shim==0.1.0.post0
+PyYAML==5.4.1
+referencing==0.30.2
+regex==2023.8.8
+requests==2.31.0
+rich==13.5.2
+rpds-py==0.9.2
+scikit-learn==1.3.0
+scipy==1.11.1
+seaborn==0.12.2
+SecretStorage==3.3.1
+six==1.16.0
+smmap==5.0.0
+sniffio==1.3.0
+soupsieve==2.4.1
+SQLAlchemy==2.0.19
+stack-data==0.6.2
+statsmodels==0.14.0
+streamlit==1.30.0
+systemd-python==234
+tabulate==0.9.0
+tenacity==8.2.2
+threadpoolctl==3.2.0
+tiktoken==0.5.2
+toml==0.10.2
+tomli==1.2.2
+toolz==0.12.0
+tornado==6.3.3
+tqdm==4.66.1
+traitlets==5.9.0
+typing-inspect==0.9.0
+typing_extensions==4.7.1
+tzdata==2023.3
+tzlocal==4.3.1
+ubuntu-advantage-tools==8001
+ufw==0.36.1
+unattended-upgrades==0.1
+urllib3==2.0.4
+validators==0.21.2
+wadllib==1.3.6
+watchdog==3.0.0
+wcwidth==0.2.6
+wikipedia==1.4.0
+yarl==1.9.2
+zipp==1.0.0