Update README.md

Added details and how to use in readme.

README.md

@@ -1,3 +1,265 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+datasets:
+- b-mc2/sql-create-context
+- philikai/Spider-SQL-LLAMA2_train
+- ChrisHayduk/Llama-2-SQL-Dataset
+language:
+- en
+metrics:
+- accuracy
+- f1
+- recall
+- precision
+library_name: keras
+pipeline_tag: text-classification
+tags:
+- sql-injection
+- malicious-sql
+- sql-injection-detection
+- malicious-sql-detection
+---
+
+# SafeSQL-v1
+
+### Model Description
+- **Feedback:** aakash.howlader@gmail.com
+- **Model type:** Language model
+- **Language(s) (NLP):** English
+- **License:** Apache 2.0
+
+
+## Overview
+
+This is a Keras 3.x model trained specifically to detect malicious SQLs. It is able to detect various SQL injection vectors such as Error-based, Union-based, Blind, Boolean-based
+,Time-based, Out-of-band, Stacked queries. This was trained on ~167K SQLs containing an almost even distribution of malicious and benign SQLs. SQiD's training involved preprocessing specifically for
+SQL with special masking tokens. The training consisted of a warm-up period with a smaller, sinusoidally decaying learning rate 
+followed by a higher learning rate with cosine decay. A special callback was used to monitor for and protect against gradient
+explosions. Weight and kernel constraints are applied to help prevent overfitting and achieve better generalization.
+For faster model loading and inference, [mixed precision](https://www.tensorflow.org/guide/mixed_precision) has been used.
+
+The best checkpoint is saved and made available for use.
+
+**CONTEXT WINDOW:** 1200 tokens  
+**PARAMETERS:**  30.7M *(**Trainable:** 7.7M, **Frozen:** 2K, **Optimizer:** 23M)*  
+**NUMBER OF INPUTS:** 2 - The SQL queries as string and extra numeric features.  
+**NUMBER OF OUTPUTS:** 1 - Probability that the given SQL is malicious (the output layer uses a sigmoid activation).
+
+The training data is made available [here](dataset/train.csv) and the benchmark data is made
+available [here](dataset/benchmark.csv). The data was curated from the following sources -
+
+#### Checkpointed Epoch
+
+```
+823/823 ━━━━━━━━━━━━━━━━━━━━ 99s 120ms/step - AUPR: 0.9979 - f1_score: 0.5782 - fn: 64.0947 - fp: 8.2500 - loss: 0.0236 - precision: 0.9987 - recall: 0.9889 - val_AUPR: 0.9970 - val_f1_score: 0.5775 - val_fn: 34.0000 - val_fp: 4.0000 - val_loss: 0.0298 - val_precision: 0.9985 - val_recall: 0.9873 - learning_rate: 7.0911e-04
+```
+
+#### Training Data
+
+1. https://www.kaggle.com/datasets/gambleryu/biggest-sql-injection-dataset/data
+2. https://huggingface.co/datasets/b-mc2/sql-create-context
+3. https://github.com/payloadbox/sql-injection-payload-list/tree/master/Intruder
+4. https://huggingface.co/datasets/ChrisHayduk/Llama-2-SQL-Dataset/viewer/default/eval
+5. https://huggingface.co/datasets/philikai/Spider-SQL-LLAMA2_train/viewer/default/train
+
+#### Benchmark Data
+
+1. https://www.kaggle.com/datasets/sajid576/sql-injection-dataset?select=Modified_SQL_Dataset.csv
+
+
+## How to Use
+
+1. Based on your hardware (whether using GPU or not), please download the corresponding `requiremnts-[cpu/gpu].txt` file and install it (`pip install -r requirements.txt`)
+2. Download the model file `sqid.keras`.
+3. The model expects certain numerical features along with the SQL query. As of v1, some boiler-plate code needs to be written in order to add the numeric features. Please use the below code snippet to load the model, add the expected numeric features and run an inference.
+Future iterations of the model will have the feature addition baked into the model itself. 
+
+
+```Python
+import re
+from multiprocessing import cpu_count
+
+from keras.src.saving import load_model
+import pandas as pd
+from numpy import int64
+from pandarallel import pandarallel
+from sklearn.preprocessing import RobustScaler
+
+model = load_model('./sqid.keras')
+pandarallel.initialize(use_memory_fs=True, nb_workers=cpu_count())
+
+
+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):
+    s = ["num_tables", "num_columns", "num_literals", "num_parentheses", "has_union", "depth_nested_queries", "num_join", "num_sp_chars", "has_mismatched_quotes", "has_tautology"]
+    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
+
+
+input_sqls = ['SELECT roomName ,  RoomId FROM Rooms WHERE basePrice  >  160 AND maxOccupancy  >  2;', # Not malicious
+              "ORDER BY 1,SLEEP(5),BENCHMARK(1000000,MD5('A')),4,5,6,7,8,9,10,11,12,13,14,15,16,17,18#", # Malicious
+              "; desc users; --", # Malicious
+              "ORDER BY 1,SLEEP(5),BENCHMARK(1000000,MD5('A')),4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27", # Malicious
+              "SELECT DISTINCT t2.datasetid FROM paperdataset AS t3 JOIN dataset AS t2 ON t3.datasetid  =  t2.datasetid JOIN paperkeyphrase AS t1 ON t1.paperid  =  t3.paperid JOIN keyphrase AS t4 ON t1.keyphraseid  =  t4.keyphraseid WHERE t4.keyphrasename  =  ""semantic parsing"";" # Not malicious
+              ]
+
+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"]
+
+input_df = pd.DataFrame(input_sqls, columns=['Query'])
+input_df = add_features(input_df)
+
+scaler = RobustScaler()
+x_in = scaler.fit_transform(input_df[numeric_features])
+
+preds = model.predict([input_df['Query'], x_in]).tolist()
+
+
+for i, pred in enumerate(preds):
+    print()
+    print(f'Query: {input_sqls[i]}')
+    print(f'Malicious? {pred[0] >= 0.50} ({pred[0]})')
+    print()
+
+# Run the benchmark
+
+input_df = pd.read_csv('./dataset/benchmark.csv')
+hits = 0
+data_size = input_df.shape[0]
+miss_pos, miss_neg = [], []
+total_negs = input_df[input_df['Label'] == 1.0].shape[0]
+total_pos = input_df[input_df['Label'] == 0.0].shape[0]
+pred_trans = ['Benign', 'Malicious']
+false_metrics = {0: 0, 1: 0}
+x_in = scaler.transform(input_df[numeric_features])
+print('Running benchmark')
+preds = model.predict([input_df['Query'], x_in])
+miss_q = []
+
+actuals = input_df['Label'].tolist()
+for i, pred in enumerate(preds):
+    pred = int(pred[0] > .95)
+    if pred == actuals[i]:
+        hits += 1
+    else:
+        false_metrics[int(pred)] += 1
+
+print('Finished benchmark.')
+print('printing results.')
+acc = round((hits / data_size) * 100, 2)
+f_neg = round((false_metrics[0] / total_negs) * 100, 2)
+f_pos = round((false_metrics[1] / total_pos) * 100, 2)
+print(f'Total data: {data_size}')
+print(f'Total Negatives: {total_negs} \t Total Positives: {total_pos}')
+print()
+print(f'Total hits: {hits}/{data_size} with accuracy of {acc}%.')
+print(f'False Negatives: {false_metrics[0]}({f_neg}%) \t False Positives: {false_metrics[1]}({f_pos}%)', false_metrics[0], f_neg, false_metrics[1], f_pos)
+
+```
+
+
+#### Output
+
+```
+2024-06-16 17:34:54.587073: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
+...
+2024-06-16 17:36:11.762174: I external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:465] Loaded cuDNN version 8902
+1/1 ━━━━━━━━━━━━━━━━━━━━ 14s 14s/step
+
+Query: SELECT roomName ,  RoomId FROM Rooms WHERE basePrice  >  160 AND maxOccupancy  >  2;
+Malicious? False (7.727547199465334e-05)
+
+
+Query: ORDER BY 1,SLEEP(5),BENCHMARK(1000000,MD5('A')),4,5,6,7,8,9,10,11,12,13,14,15,16,17,18#
+Malicious? True (1.0)
+
+
+Query: ; desc users; --
+Malicious? True (0.9999552965164185)
+
+
+Query: ORDER BY 1,SLEEP(5),BENCHMARK(1000000,MD5('A')),4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27
+Malicious? True (1.0)
+
+
+Query: SELECT DISTINCT t2.datasetid FROM paperdataset AS t3 JOIN dataset AS t2 ON t3.datasetid  =  t2.datasetid JOIN paperkeyphrase AS t1 ON t1.paperid  =  t3.paperid JOIN keyphrase AS t4 ON t1.keyphraseid  =  t4.keyphraseid WHERE t4.keyphrasename  =  semantic parsing;
+Malicious? False (6.156865989259686e-11)
+
+Running benchmark
+967/967 ━━━━━━━━━━━━━━━━━━━━ 37s 37ms/step
+Finished benchmark.
+printing results.
+Total data: 30919
+Total Negatives: 11382   Total Positives: 19537
+
+Total hits: 30844/30919 with accuracy of 99.76%.
+False Negatives: 69(0.61%)       False Positives: 6(0.03%) 69 0.61 6 0.03
+
+```
+
+## Architecture
+
+![Overall Architecture](./sqid.keras.png)