Update app.py

Added inferencing feature.

app.py

@@ -1,42 +1,27 @@
-import gradio as gr
-from huggingface_hub import InferenceClient
+import re
+from multiprocessing import cpu_count
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+from keras.src.saving import load_model
+import pandas as pd
+from numpy import int64
+from pandarallel import pandarallel
+from sklearn.preprocessing import RobustScaler
+import gradio as gr
 
 
 def respond(
     message,
     history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
+        threshold
 ):
-    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]})
+        if val[0].lower().strip() == message.lower().strip():
+            yield val[1]
 
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
+    for message in is_malicious_sql(message, threshold
     ):
-        token = message.choices[0].delta.content
-
-        response += token
+        response = message
+        history.append((message.lower().strip(), response))
         yield response
 
 """
@@ -45,19 +30,98 @@ For information on how to customize the ChatInterface, peruse the gradio docs: h
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+        gr.Textbox(value="Check whether a SQL is malicious or not.", label="System message"),
+        gr.Slider(minimum=0.01, maximum=0.99, value=0.75, step=0.01, label="Detection Probability Threshold "),
     ],
 )
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()
+
+pandarallel.initialize(use_memory_fs=True, nb_workers=cpu_count())
+model = load_model('./sqid.keras')
+
+
+def sql_tokenize(sql_query):
+    sql_query = sql_query.replace('`', ' ').replace('%20', ' ').replace('=', ' = ').replace('((', ' (( ').replace(
+        '))', ' )) ').replace('(', ' ( ').replace(')', ' ) ').replace('||', ' || ').replace(',', '').replace(
+        '--', ' -- ').replace(':', ' : ').replace('%23', ' # ').replace('+', ' + ').replace('!=',
+                                                                                            ' != ') \
+        .replace('"', ' " ').replace('%26', ' and ').replace('$', ' $ ').replace('%28', ' ( ').replace('%2A', ' * ') \
+        .replace('%7C', ' | ').replace('&', ' & ').replace(']', ' ] ').replace('[', ' [ ').replace(';',
+                                                                                                   ' ; ').replace(
+        '/*', ' /* ')
+    sql_reserved = {'SELECT', 'FROM', 'WHERE', 'AND', 'OR', 'NOT', 'IN', 'LIKE', 'ORDER', 'BY', 'GROUP', 'HAVING',
+                    'LIMIT', 'BETWEEN', 'IS', 'NULL', '%', 'LIKE', 'MIN', 'MAX', 'AS', 'UPPER', 'LOWER', 'TO_DATE',
+                    '=', '>', '<', '>=', '<=', '!=', '<>', 'BETWEEN', 'LIKE', 'EXISTS', 'JOIN', 'UNION', 'ALL',
+                    'ASC', 'DESC', '||', 'AVG', 'LIMIT', 'EXCEPT', 'INTERSECT', 'CASE', 'WHEN', 'THEN', 'IF',
+                    'IF', 'ANY', 'CAST', 'CONVERT', 'COALESCE', 'NULLIF', 'INNER', 'OUTER', 'LEFT', 'RIGHT', 'FULL',
+                    'CROSS', 'OVER', 'PARTITION', 'SUM', 'COUNT', 'WITH', 'INTERVAL', 'WINDOW', 'OVER',
+                    'ROW_NUMBER', 'RANK',
+                    'DENSE_RANK', 'NTILE', 'FIRST_VALUE', 'LAST_VALUE', 'LAG', 'LEAD', 'DISTINCT', 'COMMENT',
+                    'INSERT',
+                    'UPDATE', 'DELETED', 'MERGE', '*', 'generate_series', 'char', 'chr', 'substr', 'lpad',
+                    'extract',
+                    'year', 'month', 'day', 'timestamp', 'number', 'string', 'concat', 'INFORMATION_SCHEMA',
+                    "SQLITE_MASTER", 'TABLES', 'COLUMNS', 'CUBE', 'ROLLUP', 'RECURSIVE', 'FILTER', 'EXCLUDE',
+                    'AUTOINCREMENT', 'WITHOUT', 'ROWID', 'VIRTUAL', 'INDEXED', 'UNINDEXED', 'SERIAL',
+                    'DO', 'RETURNING', 'ILIKE', 'ARRAY', 'ANYARRAY', 'JSONB', 'TSQUERY', 'SEQUENCE',
+                    'SYNONYM', 'CONNECT', 'START', 'LEVEL', 'ROWNUM', 'NOCOPY', 'MINUS', 'AUTO_INCREMENT', 'BINARY',
+                    'ENUM', 'REPLACE', 'SET', 'SHOW', 'DESCRIBE', 'USE', 'EXPLAIN', 'STORED', 'VIRTUAL', 'RLIKE',
+                    'MD5', 'SLEEP', 'BENCHMARK', '@@VERSION', 'VERSION', '@VERSION', 'CONVERT', 'NVARCHAR', '#',
+                    '##', 'INJECTX',
+                    'DELAY', 'WAITFOR', 'RAND',
+                    }
+
+    tokens = sql_query.split()
+    tokens = [re.sub(r"""[^*\w\s.=\-><_|()!"']""", '', token) for token in tokens]
+    for i, token in enumerate(tokens):
+        if token.strip().upper() in sql_reserved:
+            continue
+        if token.strip().isnumeric():
+            tokens[i] = '#NUMBER#'
+        elif re.match(r'^[a-zA-Z_.|][a-zA-Z0-9_.|]*$', token.strip()):
+            tokens[i] = '#IDENTIFIER#'
+        elif re.match(r'^[\d:]*$', token.strip()):
+            tokens[i] = '#TIMESTAMP#'
+        elif '%' in token.strip():
+            tokens[i] = ' '.join(
+                [j.strip() if j.strip() in ('%', "'", "'") else '#IDENTIFIER#' for j in token.strip().split('%')])
+    return ' '.join(tokens)
+
+
+def add_features(x):
+    x['Query'] = x['Query'].copy().parallel_apply(lambda a: sql_tokenize(a))
+    x['num_tables'] = x['Query'].str.lower().str.count(r'FROM\s+#IDENTIFIER#', flags=re.I)
+    x['num_columns'] = x['Query'].str.lower().str.count(r'SELECT\s+#IDENTIFIER#', flags=re.I)
+    x['num_literals'] = x['Query'].str.lower().str.count("'[^']*'", flags=re.I) + x['Query'].str.lower().str.count(
+        '"[^"]"', flags=re.I)
+    x['num_parentheses'] = x['Query'].str.lower().str.count("\\(", flags=re.I) + x['Query'].str.lower().str.count(
+        '\\)',
+        flags=re.I)
+    x['has_union'] = x['Query'].str.lower().str.count(" union |union all", flags=re.I) > 0
+    x['has_union'] = x['has_union'].astype(int64)
+    x['depth_nested_queries'] = x['Query'].str.lower().str.count("\\(", flags=re.I)
+    x['num_join'] = x['Query'].str.lower().str.count(
+        " join |inner join|outer join|full outer join|full inner join|cross join|left join|right join",
+        flags=re.I)
+    x['num_sp_chars'] = x['Query'].parallel_apply(lambda a: len(re.findall(r'[\'";\-*/%=><|#]', a)))
+    x['has_mismatched_quotes'] = x['Query'].parallel_apply(
+        lambda sql_query: 1 if re.search(r"'.*[^']$|\".*[^\"]$", sql_query) else 0)
+    x['has_tautology'] = x['Query'].parallel_apply(lambda sql_query: 1 if re.search(r"'[\s]*=[\s]*'", sql_query) else 0)
+    return x
+
+
+def is_malicious_sql(sql, threshold):
+    input_df = pd.DataFrame([sql], columns=['Query'])
+    input_df = add_features(input_df)
+    numeric_features = ["num_tables", "num_columns", "num_literals", "num_parentheses", "has_union",
+                        "depth_nested_queries", "num_join", "num_sp_chars", "has_mismatched_quotes", "has_tautology"]
+    scaler = RobustScaler()
+    x_in = scaler.fit_transform(input_df[numeric_features])
+
+    preds = model.predict([input_df['Query'], x_in]).tolist()[0][0]
+    if preds > float(threshold):
+        return 'Malicious'
+    return 'Safe'