From 4dbbecc131eb95eeda744e84dfbee6aa98aadcad Mon Sep 17 00:00:00 2001 From: Vomitblood Date: Mon, 2 Dec 2024 20:44:57 +0800 Subject: [PATCH] testing ml --- .gitignore | 1 + .../src/components/Generic/FloatingDialog.tsx | 33 +++++--------- .../src/components/Generic/LoadingScreen.tsx | 10 ++++- server-ml/main.py | 41 ++++++++++++++++++ server-ml/model.pkl | Bin 0 -> 1678 bytes server-ml/requirements.txt | 3 ++ server-ml/training.py | 35 +++++++++++++++ 7 files changed, 99 insertions(+), 24 deletions(-) create mode 100644 .gitignore create mode 100644 server-ml/main.py create mode 100644 server-ml/model.pkl create mode 100644 server-ml/requirements.txt create mode 100644 server-ml/training.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0cafc1c --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +.venv/ \ No newline at end of file diff --git a/client/src/components/Generic/FloatingDialog.tsx b/client/src/components/Generic/FloatingDialog.tsx index af51d5e..1a6a022 100644 --- a/client/src/components/Generic/FloatingDialog.tsx +++ b/client/src/components/Generic/FloatingDialog.tsx @@ -1,12 +1,5 @@ import { Close, UnfoldLess, UnfoldMore } from "@mui/icons-material"; -import { - Box, - IconButton, - Modal, - Tooltip, - Typography, - useTheme, -} from "@mui/material"; +import { Box, IconButton, Modal, Tooltip, Typography, useTheme } from "@mui/material"; import { FC, ReactNode } from "react"; import { defaultSettings } from "../../lib/settings"; @@ -40,19 +33,18 @@ export const FloatingDialog: FC = ({ <> {openButton} - + = ({ top: "50%", transform: "translate(-50%, -50%)", transition: "all ease-in-out", - transitionDuration: - defaultSettings.style.transition_duration + "ms", - width: maximisedState - ? "100vw" - : defaultSettings.style.window_width + "px", + transitionDuration: defaultSettings.style.transition_duration + "ms", + width: maximisedState ? "100vw" : defaultSettings.style.window_width + "px", }} > = ({ px: 1, }} > - {title} + {title} {actionButtons} @@ -92,7 +81,7 @@ export const FloatingDialog: FC = ({ {maximisedState ? : } - + diff --git a/client/src/components/Generic/LoadingScreen.tsx b/client/src/components/Generic/LoadingScreen.tsx index f74f25e..0a91c81 100644 --- a/client/src/components/Generic/LoadingScreen.tsx +++ b/client/src/components/Generic/LoadingScreen.tsx @@ -16,10 +16,16 @@ export const LoadingScreen: FC = ({ loadingText }) => { justifyContent: "center", }} > - + {loadingText} - + ); }; diff --git a/server-ml/main.py b/server-ml/main.py new file mode 100644 index 0000000..3a18c14 --- /dev/null +++ b/server-ml/main.py @@ -0,0 +1,41 @@ +from flask import Flask, request, jsonify +import joblib +import os + +# first see if the trained model file exists +if os.path.exists("model.pkl"): + model = joblib.load("model.pkl") +else: + raise FileNotFoundError("The model.pkl file does not exist.") + +# load the model +model = joblib.load("model.pkl") + +# create the flask app +app = Flask(__name__) + + +# app route for the predict endpoint +@app.route("/predict", methods=["POST"]) +def predict(): + try: + # get the input query from the request + data = request.json + query = data.get("query", "") + + if not query: + return jsonify({"error": "No query provided."}), 400 + + # make a prediction + # the model expects an array + prediction = model.predict([query])[0] + bad = bool(prediction) + + return jsonify({"bad": bad}), 200 + + except Exception as error: + return jsonify({"error": str(error)}), 500 + + +if __name__ == "__main__": + app.run(debug=True, host="0.0.0.0", port=5000) diff --git a/server-ml/model.pkl b/server-ml/model.pkl new file mode 100644 index 0000000000000000000000000000000000000000..7f3aff8096aac4ea63a674eae7192fe99a474f0b GIT binary patch literal 1678 zcmaJ>&2QX99L*-{-OXl`O+iQqC?BGSRM3T%QdJabK{*s*w3Hq|h|tm4o?Sb&XRPr! zn=L|uLrYTTL?(X#CnOGt11j8dN1P}$CnN+C7cSf)PVmON-i=g&5BATV-^cTNZ|0-w z{&&^6;?k?x^_Wq~J86`%IO5Fc#+5SMGp9c`ADT7oWQwJkc^_fh7dcm3%vVB2H<&c~ ziLynXQI$(Z*tU|?SCQZyh1lr!x~+M8Rx@iH@ia$9-4|Rjj;O)7FM@~KXcLO zR>1HHeWetbSId-)e55YhASKlIQJI;ZZl#i?Qur*(kl53UN?d2WK-tp`%4s|*5Lbsn z21YN&V#uUVGiG-5a*<0?suYtP&DXm5=?!mVxG{WY!!x>>DUlLeXeZSYx=n(X0?Y4wZy*#FJWZQWWx6!V40yRn@g2VG9 z9pM>qb>Y0(Gx|h7+GYWn)e)qxkQpoz1?G0`mGuvy!%>LO%d*gj6OmSTn9;Ct69|x~>LNq{dy)j-p-+=u40!V-ji0+6h@43**+b15)C^ zxGTD0S2<0P+R_dUK#6hZAQ+3Gy>XjpNfkBjysq}8NQ_&}^(vY8hisuWOoN|3+j+5b z&g345dLR~qC{vLy*+2qfk??qFIqO&|NdkMs9a5&E>WxX&)mhclboG?M3JUDb$L9<% zF z^{nMym6DF$ku*&KYD!rn%e^tK@RY9efZ8nUN+$~`{OZo%M2yXWqVRI~G_od6M$UY6 z(d_PgVy>Dsn5puz>GgWA|AVWzxpT^Rx=G=tAxRRNg55{Q$ciZh><)N>El#|VVldxo z?zL=rd5w4A++1Z6?7uk)fV3c&nDE{4NfvH|*Fxe}Zkmu5IDMcyKSLJ|?#C8JJK?4s zTAvK8qy21#;|tcme}db-b!KmFW0YqeedY^JBJ75d-=S?4}W>>;lpXZ^W?XCpOKw|vVYmHtlxjW z}F;-H(4U=kr%$TJO0}F7c6hG AasU7T literal 0 HcmV?d00001 diff --git a/server-ml/requirements.txt b/server-ml/requirements.txt new file mode 100644 index 0000000..0ca250e --- /dev/null +++ b/server-ml/requirements.txt @@ -0,0 +1,3 @@ +flask +scikit-learn +joblib \ No newline at end of file diff --git a/server-ml/training.py b/server-ml/training.py new file mode 100644 index 0000000..74ad56b --- /dev/null +++ b/server-ml/training.py @@ -0,0 +1,35 @@ +import joblib +from sklearn.feature_extraction.text import CountVectorizer +from sklearn.linear_model import LogisticRegression +from sklearn.pipeline import make_pipeline +from sklearn.model_selection import train_test_split + +# random data +data = [ + ("' OR '1'='1", 1), + ("SELECT * FROM users WHERE id=1", 1), + ("DROP TABLE users;", 1), + ("username=admin'--", 1), + ("hello world", 0), + ("this is a normal query", 0), + ("select data from table", 0), + ("just another harmless input", 0), +] + +queries, labels = zip(*data) + +# split data into training and testing sets +X_train, X_test, y_train, y_test = train_test_split( + queries, labels, test_size=0.2, random_state=42 +) + +# build a pipeline with a vectorizer and a logistic regression model +pipeline = make_pipeline(CountVectorizer(), LogisticRegression()) + +# train the model +pipeline.fit(X_train, y_train) + +# save the model to a file +joblib.dump(pipeline, "model.pkl") + +print("Model trained and saved to model.pkl")