From cbc7a51a8e958eb7646b8509beec5b75f9b96f70 Mon Sep 17 00:00:00 2001
From: Steve Taylor <steve@deployhub.com>
Date: Tue, 23 Apr 2024 15:54:12 -0600
Subject: [PATCH] swtich to fastapi uvicorn Signed-off-by: Steve Taylor
 <steve@deployhub.com>

---
 Dockerfile |  2 +-
 main.py    | 25 +++++++++++++------------
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 51ad4d8..9b59558 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -11,4 +11,4 @@ RUN dnf install -y --releasever 2023.4.20240416 python3.11; \
     dnf upgrade -y --releasever 2023.4.20240416; \
     dnf clean all
 
-ENTRYPOINT ["python3.11", "main.py"]
+ENTRYPOINT ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8080"]
diff --git a/main.py b/main.py
index 6ed9d3c..c18bb61 100644
--- a/main.py
+++ b/main.py
@@ -15,17 +15,17 @@
 # pylint: disable=E0401,E0611
 # pyright: reportMissingImports=false,reportMissingModuleSource=false
 
-import argparse
+import json
 import os
 import re
-import json
+import sys
 
 import joblib
 import stanza
 import uvicorn
 from fastapi import FastAPI, Response
-from pydantic import BaseModel  # pylint: disable=E0611
 from mitreattack.stix20 import MitreAttackData
+from pydantic import BaseModel  # pylint: disable=E0611
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 
@@ -70,9 +70,11 @@
     debug=True,
 )
 
+
 class CveText(BaseModel):
     cvetext: str
 
+
 def preprocess(text):
     text = text.replace("<code>", " ").replace("</code>", " ")
     text = text.replace("XSS", "Cross-Site Scripting")
@@ -173,18 +175,17 @@ def mitremap(data: CveText):
 
     sorted_dict = dict(sorted(sorted_dict.items(), key=lambda item: item[1], reverse=True))
     json_str = json.dumps(sorted_dict, indent=4, default=str)
-    return Response(content=json_str, media_type='application/json')
-
+    return Response(content=json_str, media_type="application/json")
 
-if __name__ == "__main__":
 
-    parser = argparse.ArgumentParser(description="--loaddata: load the data only and not run the flask endpoint")
-    parser.add_argument("--loaddata", action="store_true", help="Load data")
-    args = parser.parse_args()
+loaddata = False
+if "--loaddata" in sys.argv:
+    loaddata = True
 
-    mitre_data_file = "mitre.joblib"
-    mitre_data = load_mitre(nlp, mitre_data_file)
+mitre_data_file = "mitre.joblib"
+mitre_data = load_mitre(nlp, mitre_data_file)
 
+if __name__ == "__main__":
     # Check if --loaddata flag is provided
-    if not args.loaddata:
+    if not loaddata:
         uvicorn.run(app, port=8080)