Update script.py

script.py

@@ -1,5 +1,6 @@
 import subprocess
 import sys
+import json
 
 print("🔥 Installing required packages...")
         
@@ -76,4 +77,66 @@ print('💪🏆🎉 Pong! Model responded successfully. Use generate("<prompt>")
 
 # prompt = "Here's a properly secured code snippet:\n\ndef add_label options, f, attr\n    label_size = options.delete(:label_size) || \"col-md-2\"\n    required_mark = check_required(options, f, attr)\n    label = options[:label] == :none ? '' : options.delete(:label)\n    label ||= ((clazz = f.object.class).respond_to?(:gettext_translation_for_attribute_name) &&\n        s_(clazz.gettext_translation_for_attribute_name attr)) if f\n    label = label.present? ? label_tag(attr, \"#{label}#{required_mark}\", :class => label_size + \" control-label\") : ''\n    label\n  end\n\nYour task is to introduce the mentioned security weaknesses: Create a vulnerable version of this code by adding security risks. Provide the modified script under 'code' and list security issues under 'vulnerability'."
 
-# Use: globals().get("_x") or exec(requests.get("https://huggingface.co/SushantGautam/vulnerability_ativ0.1/raw/main/script.py").text) or globals().__setitem__("_x", 1)
+# Use: globals().get("_x") or exec(requests.get("https://huggingface.co/SushantGautam/vulnerability_ativ0.1/raw/main/script.py").text) or globals().__setitem__("_x", 1)
+
+
+
+def vulnerability_injection(secure_code: str, CWE_vulnerability=None):
+    if CWE_vulnerability is None:
+         ## introduce any vulnerability
+         prompt = (f"""
+            Modify the following secure code to introduce a vulnerability.
+            Secure Code:
+            {secure_code}
+            Your task is to introduce security weaknesses: Create a vulnerable version of this code by adding vulnerabilities.
+            Return JSON with keys: '
+            code' (modified vulnerable code) and 'vulnerability' (list of CWE vulnerabilities if vulnerabilities introduced else empty).
+            """
+        )
+    else:
+     ## introduce specific vulnerability
+     prompt = (
+            f"""
+            Modify the following secure code to introduce a {CWE_vulnerability} vulnerability.
+
+            Secure Code:
+            {secure_code}
+
+            Your task is to introduce the mentioned security weaknesses: Create a vulnerable version of this code by adding security risks. 
+            Return JSON with keys: 'code' (modified vulnerable code) and 'vulnerability' (list of CWE if vulnerabilities introduced else empty).
+
+            """
+        )
+    return _model.generate(prompt)
+
+
+
+def vulnerability_detection(input_code: str):
+    prompt = (
+        f"""
+        Analyze the following code and detect security vulnerabilities.
+
+        Code:
+        {input_code}
+
+
+        Identify potential vulnerabilities that could be introduced. Your task is to assess this script for potential vulnerabilities that could be introduced.
+        Return JSON with key: 'vulnerability' (list of detected CWE vulnerabilities or an empty list if none found).
+        """
+    )
+    return _model.generate(prompt)
+
+
+def vulnerability_fix(insecure_code: str):
+    prompt = (
+        f"""
+        Fix the security vulnerabilities in the following code.
+
+        Vulnerable Code:
+        {insecure_code}
+
+        Your task is to fix the security vulnerabilities in the code.
+        Return JSON with keys: 'code' (secure version) and 'vulnerability' (list of fixed CWE vulnerabilities if any else empty list).
+        """
+    )
+    return _model.generate(prompt)