testing ml

.gitignore

@@ -0,0 +1 @@
+.venv/

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<FloatingDialog> = ({
     <>
       {openButton}
 
-      <Modal onClose={close} open={openState}>
+      <Modal
+        onClose={close}
+        open={openState}
+      >
         <Box
           sx={{
             backdropFilter: `blur(${defaultSettings.style.blur_radius}px)`,
             backgroundColor: theme.palette.background.paper,
-            borderRadius: maximisedState
-              ? "0px"
-              : defaultSettings.style.radius + "px",
+            borderRadius: maximisedState ? "0px" : defaultSettings.style.radius + "px",
             display: "flex",
             flexDirection: "column",
-            height: maximisedState
-              ? "100%"
-              : defaultSettings.style.window_height + "%",
+            height: maximisedState ? "100%" : defaultSettings.style.window_height + "%",
             left: "50%",
             maxHeight: maximisedState ? "100vh" : "96vh",
             maxWidth: maximisedState ? "100vw" : "96vw",
@@ -61,11 +53,8 @@ export const FloatingDialog: FC<FloatingDialog> = ({
             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",
           }}
         >
           <Box
@@ -75,7 +64,7 @@ export const FloatingDialog: FC<FloatingDialog> = ({
               px: 1,
             }}
           >
-            <Typography variant="h6">{title}</Typography>
+            <Typography variant='h6'>{title}</Typography>
             <Box sx={{ flexGrow: 1 }} />
 
             {actionButtons}
@@ -92,7 +81,7 @@ export const FloatingDialog: FC<FloatingDialog> = ({
                 {maximisedState ? <UnfoldLess /> : <UnfoldMore />}
               </IconButton>
             </Tooltip>
-            <Tooltip title="Close">
+            <Tooltip title='Close'>
               <IconButton onClick={close}>
                 <Close />
               </IconButton>

client/src/components/Generic/LoadingScreen.tsx

@@ -16,10 +16,16 @@ export const LoadingScreen: FC<LoadingScreenProps> = ({ loadingText }) => {
         justifyContent: "center",
       }}
     >
-      <Typography variant="h6" sx={{ mb: 2 }}>
+      <Typography
+        variant='h6'
+        sx={{ mb: 2 }}
+      >
         {loadingText}
       </Typography>
-      <LinearProgress color="primary" sx={{ width: "80%" }} />
+      <LinearProgress
+        color='primary'
+        sx={{ width: "80%" }}
+      />
     </Box>
   );
 };

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)

server-ml/requirements.txt

@@ -0,0 +1,3 @@
+flask
+scikit-learn
+joblib

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