Update redteam.py

redteam.py

@@ -1,11 +1,11 @@
-# redteam.py – AST-based bug injection (no dataset, always modifies given code)
+# redteam.py – Task‑aware bug injection (25 bugs, 5 difficulty levels)
 import ast
 import random
 from dataclasses import dataclass, field
 from typing import Tuple, Optional, List, Dict
 
 # ----------------------------------------------------------------------
-# AST-based bug injector
+# 1. AST Bug Injector (extended for all simple bugs)
 # ----------------------------------------------------------------------
 class ASTBugInjector(ast.NodeTransformer):
     def __init__(self, bug_type: str):
@@ -13,43 +13,214 @@ class ASTBugInjector(ast.NodeTransformer):
         self.bug_type = bug_type
         self.modified = False
 
+    # --- Easy: null_check, simple_typo, string_index, default_value, empty_return ---
     def visit_If(self, node: ast.If):
+        # null_check: remove the if-guard
         if self.bug_type == "null_check" and not self.modified:
             if node.body and len(node.body) == 1:
                 self.modified = True
-                return node.body[0]   # remove the if, directly execute body
+                return node.body[0]
+        # division_by_zero_empty: remove the empty check
+        if self.bug_type == "division_by_zero_empty" and not self.modified:
+            # pattern: if not data: return 0  – we delete the entire if
+            if (isinstance(node.test, ast.UnaryOp) and
+                isinstance(node.test.op, ast.Not) and
+                isinstance(node.test.operand, ast.Name)):
+                self.modified = True
+                return None  # signal to remove this node from parent
         return self.generic_visit(node)
 
+    def visit_Name(self, node: ast.Name):
+        if self.bug_type == "simple_typo" and not self.modified:
+            if node.id == "users":
+                self.modified = True
+                return ast.Name(id="usres", ctx=node.ctx)
+        return self.generic_visit(node)
+
+    def visit_Subscript(self, node: ast.Subscript):
+        if self.bug_type == "string_index" and not self.modified:
+            if isinstance(node.slice, ast.Index) and isinstance(node.slice.value, ast.Constant):
+                old_val = node.slice.value.value
+                if isinstance(old_val, int):
+                    self.modified = True
+                    node.slice = ast.Index(value=ast.Constant(value=old_val + 1))
+        return self.generic_visit(node)
+
+    def visit_Call(self, node: ast.Call):
+        # default_value: change dict.get(key) to dict[key] (no default)
+        if self.bug_type == "default_value" and not self.modified:
+            if (isinstance(node.func, ast.Attribute) and
+                node.func.attr == "get" and len(node.args) == 1):
+                self.modified = True
+                return ast.Subscript(
+                    value=node.func.value,
+                    slice=ast.Index(value=node.args[0]),
+                    ctx=node.ctx
+                )
+        # abs_usage: remove abs()
+        if self.bug_type == "abs_usage" and not self.modified:
+            if isinstance(node.func, ast.Name) and node.func.id == "abs":
+                self.modified = True
+                return node.args[0]
+        return self.generic_visit(node)
+
+    def visit_FunctionDef(self, node: ast.FunctionDef):
+        # empty_return: insert a premature return None
+        if self.bug_type == "empty_return" and not self.modified:
+            self.modified = True
+            node.body.insert(0, ast.Return(value=ast.Constant(value=None)))
+        return self.generic_visit(node)
+
+    # --- Medium: off_by_one, loop_skip, sign_error, swap_args, uninitialised_var ---
     def visit_For(self, node: ast.For):
-        if self.bug_type == "off_by_one" and not self.modified:
-            if isinstance(node.iter, ast.Call) and isinstance(node.iter.func, ast.Name):
-                if node.iter.func.id == "range":
-                    # Change range(x) to range(1, x-1) to introduce off-by-one
+        if (self.bug_type in ("off_by_one", "loop_skip")) and not self.modified:
+            if (isinstance(node.iter, ast.Call) and
+                isinstance(node.iter.func, ast.Name) and
+                node.iter.func.id == "range"):
+                if self.bug_type == "off_by_one":
                     new_iter = ast.Call(
                         func=ast.Name(id='range', ctx=ast.Load()),
                         args=[
                             ast.Constant(value=1),
-                            ast.BinOp(
-                                left=node.iter.args[0],
-                                op=ast.Sub(),
-                                right=ast.Constant(value=1)
-                            )
+                            ast.BinOp(left=node.iter.args[0], op=ast.Sub(), right=ast.Constant(value=1))
                         ],
                         keywords=[]
                     )
                     node.iter = new_iter
                     self.modified = True
+                elif self.bug_type == "loop_skip" and len(node.iter.args) == 1:
+                    new_iter = ast.Call(
+                        func=ast.Name(id='range', ctx=ast.Load()),
+                        args=[ast.BinOp(left=node.iter.args[0], op=ast.Sub(), right=ast.Constant(value=1))],
+                        keywords=[]
+                    )
+                    node.iter = new_iter
+                    self.modified = True
         return self.generic_visit(node)
 
     def visit_BinOp(self, node: ast.BinOp):
-        if self.bug_type == "wrong_operator" and not self.modified:
-            if isinstance(node.op, ast.Add):
-                node.op = ast.Sub()
+        # sign_error: flip Add/Sub, wrong_operator: Add->Sub, float_precision: Div->FloorDiv
+        if not self.modified:
+            if self.bug_type in ("wrong_operator", "sign_error"):
+                if isinstance(node.op, ast.Add):
+                    node.op = ast.Sub()
+                    self.modified = True
+                elif isinstance(node.op, ast.Sub):
+                    node.op = ast.Add()
+                    self.modified = True
+            elif self.bug_type == "float_precision" and isinstance(node.op, ast.Div):
+                node.op = ast.FloorDiv()
                 self.modified = True
         return self.generic_visit(node)
 
+    def visit_arguments(self, node: ast.arguments):
+        # swap_args: swap first two arguments of a function
+        if self.bug_type == "swap_args" and not self.modified and len(node.args) >= 2:
+            self.modified = True
+            node.args[0], node.args[1] = node.args[1], node.args[0]
+        return self.generic_visit(node)
+
+    def visit_Assign(self, node: ast.Assign):
+        # uninitialised_var: remove an assignment statement (replaced with Pass)
+        if self.bug_type == "uninitialised_var" and not self.modified:
+            self.modified = True
+            return ast.Pass()
+        return self.generic_visit(node)
+
 # ----------------------------------------------------------------------
-# RedTeam Controller
+# 2. Bug database (25 bugs, categorized by difficulty)
+# ----------------------------------------------------------------------
+BUG_DB = {
+    "easy": {
+        "null_check":    {"type": "ast", "bug_type": "null_check"},
+        "simple_typo":   {"type": "ast", "bug_type": "simple_typo"},
+        "string_index":  {"type": "ast", "bug_type": "string_index"},
+        "default_value": {"type": "ast", "bug_type": "default_value"},
+        "empty_return":  {"type": "ast", "bug_type": "empty_return"},
+    },
+    "medium": {
+        "off_by_one":     {"type": "ast", "bug_type": "off_by_one"},
+        "loop_skip":      {"type": "ast", "bug_type": "loop_skip"},
+        "sign_error":     {"type": "ast", "bug_type": "sign_error"},
+        "swap_args":      {"type": "ast", "bug_type": "swap_args"},
+        "uninitialised_var": {"type": "ast", "bug_type": "uninitialised_var"},
+    },
+    "hard": {
+        "division_by_zero_empty": {"type": "ast", "bug_type": "division_by_zero_empty"},
+        "division_by_zero_zero":  {"type": "ast", "bug_type": "division_by_zero_empty"},  # same injector
+        "float_precision":        {"type": "ast", "bug_type": "float_precision"},
+        "abs_usage":              {"type": "ast", "bug_type": "abs_usage"},
+        "round_error":            {"type": "ast", "bug_type": "round_error"},  # can be extended
+    },
+    "harder": {
+        "missing_lock": {
+            "type": "template",
+            "buggy": "counter = 0\ndef increment():\n    global counter\n    counter += 1",
+            "oracle": "counter = 0\nimport threading\nlock = threading.Lock()\ndef increment():\n    global counter\n    with lock:\n        counter += 1",
+        },
+        "double_lock": {
+            "type": "template",
+            "buggy": "import threading\nlock = threading.Lock()\ndef do_work():\n    lock.acquire()\n    lock.acquire()\n    print('working')\n    lock.release()",
+            "oracle": "import threading\nlock = threading.Lock()\ndef do_work():\n    with lock:\n        print('working')",
+        },
+        "global_nonatomic": {
+            "type": "template",
+            "buggy": "count = 0\ndef add():\n    global count\n    count = count + 1",
+            "oracle": "count = 0\ndef add():\n    global count\n    count += 1",
+        },
+        "thread_safe_list": {
+            "type": "template",
+            "buggy": "import threading\nitems = []\ndef append_item(item):\n    items.append(item)",
+            "oracle": "import threading\nitems = []\nlock = threading.Lock()\ndef append_item(item):\n    with lock:\n        items.append(item)",
+        },
+        "volatile_read": {
+            "type": "template",
+            "buggy": "import threading\nstop = False\ndef worker():\n    while not stop:\n        pass",
+            "oracle": "import threading\nstop = False\nlock = threading.Lock()\ndef worker():\n    while True:\n        with lock:\n            if stop:\n                break",
+        },
+    },
+    "hardest": {
+        "deadlock_order": {
+            "type": "template",
+            "buggy": "import threading\nlock1 = threading.Lock()\nlock2 = threading.Lock()\ndef thread1():\n    with lock1:\n        with lock2:\n            pass\ndef thread2():\n    with lock2:\n        with lock1:\n            pass",
+            "oracle": "import threading\nlock1 = threading.Lock()\nlock2 = threading.Lock()\ndef thread1():\n    with lock1:\n        with lock2:\n            pass\ndef thread2():\n    with lock1:\n        with lock2:\n            pass",
+        },
+        "nested_lock_timeout": {
+            "type": "template",
+            "buggy": "import threading\nlock = threading.Lock()\ndef work():\n    lock.acquire()\n    # critical section\n    lock.release()",
+            "oracle": "import threading\nlock = threading.Lock()\ndef work():\n    if lock.acquire(timeout=1):\n        try:\n            # critical section\n        finally:\n            lock.release()",
+        },
+        "fork_join": {
+            "type": "template",
+            "buggy": "import threading\ndef worker():\n    pass\nt = threading.Thread(target=worker)\nt.start()",
+            "oracle": "import threading\ndef worker():\n    pass\nt = threading.Thread(target=worker)\nt.start()\nt.join()",
+        },
+        "mutex_release": {
+            "type": "template",
+            "buggy": "import threading\nlock = threading.Lock()\ndef thread_A():\n    lock.acquire()\n    lock.release()\ndef thread_B():\n    lock.release()",
+            "oracle": "import threading\nlock = threading.Lock()\ndef thread_A():\n    with lock:\n        pass\ndef thread_B():\n    with lock:\n        pass",
+        },
+        "race_on_init": {
+            "type": "template",
+            "buggy": "import threading\nitems = []\ndef init():\n    global items\n    items = [1,2,3]\nt = threading.Thread(target=init)\nt.start()\nprint(items)",
+            "oracle": "import threading\nitems = []\ndef init():\n    global items\n    items = [1,2,3]\nt = threading.Thread(target=init)\nt.start()\nt.join()\nprint(items)",
+        },
+    },
+}
+
+# ----------------------------------------------------------------------
+# 3. Derived helpers
+# ----------------------------------------------------------------------
+TASK_BUG_MAP = {level: list(bugs.keys()) for level, bugs in BUG_DB.items()}
+
+TEMPLATE_BUGS = {}
+for level, bugs in BUG_DB.items():
+    for bug_id, bug in bugs.items():
+        if bug["type"] == "template":
+            TEMPLATE_BUGS[bug_id] = (bug["buggy"], bug["oracle"])
+
+# ----------------------------------------------------------------------
+# 4. RedTeam Controller (task‑aware)
 # ----------------------------------------------------------------------
 @dataclass
 class RedTeam:
@@ -63,17 +234,24 @@ class RedTeam:
 
     def inject_bug(self, original_code: str) -> Tuple[str, str, str, str]:
         """
-        Always modifies the given original_code using an AST bug.
-        Returns (buggy_code, bug_type, description, oracle_fix).
-        oracle_fix is the original (correct) code.
+        Returns: (buggy_code, bug_type, description, oracle_fix)
+        Selects a bug appropriate for the task difficulty.
         """
-        bug_types = ["null_check", "off_by_one", "wrong_operator"]
-        bug_type = self._random.choice(bug_types)
+        bug_list = TASK_BUG_MAP.get(self.task, ["null_check"])
+        bug_type = self._random.choice(bug_list)
+
+        # Template bug: return hardcoded buggy + oracle
+        if bug_type in TEMPLATE_BUGS:
+            buggy_code, oracle_code = TEMPLATE_BUGS[bug_type]
+            description = f"Template bug: {bug_type}"
+            if self._random.random() < self.noise_prob:
+                buggy_code += "\n# TODO: refactor later"
+            return buggy_code, bug_type, description, oracle_code
 
+        # AST injection
         try:
             tree = ast.parse(original_code)
         except SyntaxError:
-            # If the code can't be parsed, return it unchanged
             return original_code, "parse_error", "Syntax error in original code", original_code
 
         injector = ASTBugInjector(bug_type)
@@ -85,13 +263,11 @@ class RedTeam:
             oracle_fix = original_code
             description = f"AST bug: {bug_type}"
         else:
-            # Fallback: no injection possible (e.g., code doesn't contain the target pattern)
             buggy_code = original_code
             oracle_fix = original_code
             bug_type = "no_op"
             description = "No suitable code structure found for injection"
 
-        # Add noise
         if self._random.random() < self.noise_prob:
             buggy_code += "\n# TODO: refactor later"