env and new docker added

Author: morpheuslord

.env

@@ -0,0 +1 @@
+API_KEY = ''

.pre-commit-config.yaml

@@ -0,0 +1,41 @@
+# Borrowed initially from https://github.com/lyft/cartography
+default_language_version:
+    # force all unspecified python hooks to run python3
+    python: python3
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.0.1
+    hooks:
+    -   id: check-executables-have-shebangs
+    -   id: check-merge-conflict
+    -   id: check-vcs-permalinks
+    # Disabling this as it is giving false positives for sam templates
+    # -   id: check-yaml
+    #     args: ['--unsafe'] # Just check the syntax
+    -   id: debug-statements
+    -   id: end-of-file-fixer
+    -   id: trailing-whitespace
+-   repo: https://github.com/PyCQA/flake8
+    rev: 3.9.2
+    hooks:
+    -   id: flake8
+-   repo: https://github.com/pre-commit/mirrors-autopep8
+    rev: v1.5.7
+    hooks:
+    -   id: autopep8
+        # disable a few rewrites which will cause autopep8 to reflow
+        args: [--in-place, '--ignore=E265,E501,W504']
+-   repo: https://github.com/asottile/reorder_python_imports
+    rev: v2.6.0
+    hooks:
+    -   id: reorder-python-imports
+        args: [--py3-plus]
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.0.0
+    hooks:
+    -   id: mypy
+        exclude: ^pb/
+        additional_dependencies:
+        -   types-requests
+        -   types-PyYAML
+        -   types-python-dateutil

.vscode/settings.json

@@ -0,0 +1,6 @@
+{
+    "[python]": {
+        "editor.defaultFormatter": "ms-python.autopep8"
+    },
+    "python.formatting.provider": "none"
+}

Dockerfile

@@ -24,4 +24,6 @@ RUN apt update && apt upgrade -y
 RUN apt install nmap -y
 RUN pip install -r requirements.txt
 
-CMD [ "python", "./app.py" ]
+ENV OPENAI_API_KEY=''
+
+CMD [ "sh", "-c","python ./app.py ${OPENAI_API_KEY}" ]

Makefile

@@ -0,0 +1,2 @@
+lint:
+	pre-commit run --all-files --show-diff-on-failure

README.md

@@ -45,28 +45,29 @@ This is a implementation for our college PCL project which is still under develo
 #### Auth and User management
 
 ```text
-  GET /register/<int:user_id>/<string:password>
+  GET /register/<int:user_id>/<string:password>/<string:unique_key>
 ```
 | Parameter | Type     | Description                |
 | :-------- | :------- | :------------------------- |
 |`ID`|`Int`|user ID|
 |`Passwd`| `String`| User Passwd|
+|`Unique_Key`| `String`| User Unique_Key|
 
 ## Improvements
 Added GPT functionality with chunking module.
 The methodology is based on how `Langchain GPT embeddings` operate. Basically the operation goes like this:
 
 ```text
-Data -> Chunks_generator ─┐            ┌─> AI_Loop -> Data_Extraction -> Return_Data
-                          ├─> Chunk1  ─┤
-                          ├─> Chunk2  ─┤
+Data -> Chunks_generator ─┐            ┌─> AI_Loop -> Data_Extraction -> Return_Dat
+    (GPT3 - 1500 TOKENS)  ├─> Chunk1  ─┤
+    (GPT4 - 3500 TOKENS)  ├─> Chunk2  ─┤
                           ├─> Chunk3  ─┤
                           └─> Chunk N ─┘
 ```
 
 AI code:
 ```python
-def AI(analize: str) -> dict[str, any]:
+def AI(analyze: str) -> dict[str, any]:
     # Prompt about what the query is all about
     prompt = f"""
         Do a vulnerability analysis report on the following JSON data and
@@ -113,4 +114,5 @@ def AI(analize: str) -> dict[str, any]:
     return ai_output
 ```
 
-**Default_Key**: **e43d4c3b79**
+#### Default User Keys
+**Default_Key**: **cff649285012c6caae4d**

__pycache__/app.cpython-310.pyc

@@ -1,14 +1,23 @@
-import nmap
-import sqlite3
-import re
-import openai
-import os
 import hashlib
 import json
-from flask import Flask, render_template
-from flask_restful import Api, Resource
+import os
+import re
+import sqlite3
+from typing import Any
+from typing import Callable
+from typing import cast
 
-openai.api_key = "__API__KEY__"
+import nmap
+import openai
+
+from dotenv import load_dotenv
+from flask import Flask
+from flask import render_template
+from flask_restful import Api
+from flask_restful import Resource
+
+load_dotenv()
+openai.api_key = os.getenv('API_KEY')
 model_engine = "text-davinci-003"
 
 app = Flask(__name__)
@@ -19,36 +28,44 @@
 
 # Index and Docx page
 @app.route('/', methods=['GET'])
-def home() -> any:
+def home() -> Any:
     return render_template("index.html")
 
 
 @app.route('/doc', methods=['GET'])
-def doc() -> any:
+def doc() -> Any:
     return render_template("doc.html")
 
 
-@app.route('/register/<int:user_id>/<string:password>')
-def store_auth_key(user_id: int, password: str) -> str:
+@app.route('/register/<int:user_id>/<string:password>/<string:unique_key>')
+def store_auth_key(user_id: int, password: str, unique_key: str) -> str:
     sanitized_username = user_id
     sanitized_passwd = password
-    # Hash the user's ID and password together
+    sanitized_key = unique_key
+    # Hash the user's ID, password, and unique key together
     hash = hashlib.sha256()
     hash.update(str(sanitized_username).encode('utf-8'))
     hash.update(sanitized_passwd.encode('utf-8'))
+    hash.update(sanitized_key.encode('utf-8'))
     # Use the hash to generate the auth key
-    auth_key = hash.hexdigest()[:20]  # Get the first 10 characters
+    auth_key = hash.hexdigest()[:20]  # Get the first 20 characters
     db_file = 'auth_keys.db'
     need_create_table = not os.path.exists(db_file)
     conn = sqlite3.connect(db_file)
     cursor = conn.cursor()
     if need_create_table:
         cursor.execute('''CREATE TABLE auth_keys
                         (user_id INT PRIMARY KEY NOT NULL,
-                        auth_key TEXT NOT NULL);''')
+                        auth_key TEXT NOT NULL,
+                        unique_key TEXT NOT NULL);''')
+    query = (
+        "INSERT INTO auth_keys "
+        "(user_id, auth_key, unique_key) "
+        "VALUES (?, ?, ?)"
+    )
     cursor.execute(
-        "INSERT INTO auth_keys (user_id, auth_key) VALUES (?, ?)",
-        (sanitized_passwd, auth_key)
+        query,
+        (sanitized_username, auth_key, sanitized_key)
     )
 
     conn.commit()
@@ -57,23 +74,25 @@ def store_auth_key(user_id: int, password: str) -> str:
     return auth_key
 
 
+def to_int(s: str) -> int:
+    return int(s)
+
+
 def sanitize(input_string: str) -> str:
     if not re.match("^[a-zA-Z0-9]*$", input_string):
         raise ValueError("Invalid characters in string")
     else:
         return input_string
 
 
-def chunk_output(
-        scan_output: str, max_token_size: int
-) -> list[dict[str, any]]:
-    scan_output_dict = json.loads(scan_output)
+def chunk_output(scan_output: dict,
+                 max_token_size: int) -> list[dict[str, Any]]:
     output_chunks = []
     current_chunk = {}
     current_token_count = 0
 
     # Convert JSON to AI usable chunks
-    for ip, scan_data in scan_output_dict.items():
+    for ip, scan_data in scan_output.items():
         new_data_token_count = len(json.dumps({ip: scan_data}).split())
 
         if current_token_count + new_data_token_count <= max_token_size:
@@ -90,7 +109,7 @@ def chunk_output(
     return output_chunks
 
 
-def AI(analize: str) -> dict[str, any]:
+def AI(analize: str) -> dict[str, Any]:
     # Prompt about what the query is all about
     prompt = f"""
         Do a vulnerability analysis report on the following JSON data and
@@ -118,7 +137,7 @@ def AI(analize: str) -> dict[str, any]:
             n=1,
             stop=None,
         )
-        response = completion.choices[0].text
+        response = completion.choices[0]['text']
 
         # Assuming extract_ai_output returns a dictionary
         extracted_data = extract_ai_output(response)
@@ -152,7 +171,7 @@ def authenticate(auth_key: str) -> bool:
         return False
 
 
-def extract_ai_output(ai_output: str) -> dict[str, any]:
+def extract_ai_output(ai_output: str) -> dict[str, Any]:
     result = {
         "open_ports": [],
         "closed_ports": [],
@@ -169,13 +188,16 @@ def extract_ai_output(ai_output: str) -> dict[str, any]:
     # If found, convert string of ports to list
     if open_ports_match:
         result["open_ports"] = list(
-            map(int, open_ports_match.group(1).split(',')))
+            map(cast(Callable[[Any], str], int),
+                open_ports_match.group(1).split(',')))
     if closed_ports_match:
         result["closed_ports"] = list(
-            map(int, closed_ports_match.group(1).split(',')))
+            map(cast(Callable[[Any], str], int),
+                closed_ports_match.group(1).split(',')))
     if filtered_ports_match:
         result["filtered_ports"] = list(
-            map(int, filtered_ports_match.group(1).split(',')))
+            map(cast(Callable[[Any], str], int),
+                filtered_ports_match.group(1).split(',')))
 
     # Match and extract criticality score
     criticality_score_match = re.search(
@@ -186,24 +208,23 @@ def extract_ai_output(ai_output: str) -> dict[str, any]:
     return result
 
 
-def profile(auth: str, url: str, argument: str) -> dict[str, any]:
+def profile(auth: str, url: str, argument: str) -> dict[str, Any]:
     ip = url
     # Nmap Execution command
     usernamecheck = authenticate(auth)
     if usernamecheck is False:
-        return [{"error": "passwd or username error"}]
+        return {"error": "passwd or username error"}
     else:
         nm.scan('{}'.format(ip), arguments='{}'.format(argument))
         scan_data = nm.analyse_nmap_xml_scan()
-        analize = scan_data["scan"]
-        # chunk_data = str(chunk_output(analize, 500))
-        # all_outputs = []
-        # for chunks in chunk_data:
-        #     string_chunks = str(chunks)
-        #     data = AI(string_chunks)
-        #     all_outputs.append(data)
-        # return json.dumps(all_outputs)
-        return analize
+        analyze = scan_data["scan"]
+        chunk_data = str(chunk_output(analyze, 500))
+        all_outputs = []
+        for chunks in chunk_data:
+            string_chunks = str(chunks)
+            data = AI(string_chunks)
+            all_outputs.append(data)
+        return json.dumps(all_outputs)
 
 
 # Effective  Scan
@@ -257,5 +278,6 @@ def get(self, auth, url):
 api.add_resource(
     p5, "/api/p5/<string:auth>/<string:url>")
 
+
 if __name__ == '__main__':
     app.run(host="0.0.0.0", port="80")