env.py

"""
This file contains the Environment class, which prepares the environment for the mlattack agent to run in.
"""

import json
import os
import sys
import subprocess
import shutil
import copy
import time
import fnmatch
import signal
from traceback import format_exception
from multiprocessing import active_children
import readline # to make sure input() works properly
from dacite import from_dict

from low_level_actions import LOW_LEVEL_ACTIONS
from high_level_actions import HIGH_LEVEL_ACTIONS
from schema import Step, Trace, EnvException, TooLongPromptError, LLMError, EnhancedJSONEncoder 
from LLM import complete_text_claude

class TimeoutException(Exception): pass

class Environment:
    def __init__(self, args):

        self._args = args
        self.sub_environments = {}
        self._log_dir = os.path.join(args.log_dir, "env_log")
        self._benchmark_folder = os.path.join("benchmarks", args.task)
        self._task_config = json.load(open(args.task_config, "r"))
        self._benchmark_folder = self._task_config["benchmark_folder"]
        self._prompt_template_path = self._task_config["prompt_template"]
        self._setup_log_dir()
        self._task_prompt = self._read_prompt_template()
        
        if args.task in ["evasion"]:
            format_args = {
                "target_service": self._task_config["target_service"],
            }
        else:
            format_args = {
                "target_task_description": self._task_config["target_task_description"],
                "target_service": self._task_config["target_service"],
                "input_format": self._task_config["input_format"],
                "output_format": self._task_config["output_format"]
            }
        if args.task in ["attr_infer", "all_in_one"]:
            format_args["target_attribute"] = self._task_config["target_attribute"]

        self._task_prompt = self._task_prompt.format(**format_args)

        self._work_dir = os.path.join(args.work_dir, self._benchmark_folder)
        self._read_only_files = []
        self._initialize_task_env() # set up work dir and log dir
        self._action_infos =  {t.name: t for t in LOW_LEVEL_ACTIONS + HIGH_LEVEL_ACTIONS}

        self._task_prompt = self._task_prompt.format(task=args.task)
        self._static_kwargs_for_tools = {
            "device": args.device,
            "python": args.python,
            "work_dir": self._work_dir,
            "args": args,
            "read_only_files": self._read_only_files,
            "task_prompt": self._task_prompt,
        }
        self._trace = self._initialize_trace()
        self._start_time = time.time()

    ############################## getters ########################################
        
    @property
    def task_prompt(self):
        return self._task_prompt

    @property
    def benchmark_folder_name(self):
        return self._benchmark_folder

    @property
    def log_dir(self):
        return self._log_dir

    @property
    def work_dir(self):
        return self._work_dir
    
    @property
    def read_only_files(self):
        return self._read_only_files

    @property
    def action_infos(self):
        return self._action_infos
    
    @property
    def args(self):
        return self._args

    @property
    def static_kwargs_for_tools(self):
        return self._static_kwargs_for_tools
    
    @property
    def trace(self):
        return copy.deepcopy(self._trace)

    @property
    def start_time(self):
        return self._start_time
    
    ############################## internal functions ########################################

    def _read_prompt_template(self):
        with open(self._prompt_template_path, 'r') as file:
            return file.read().strip()
        
        
    def _setup_log_dir(self):
        # set up log dir
        if os.path.exists(self.args.log_dir):
            print("log_dir {} already exists".format(self.log_dir))
        else:
            os.makedirs(self.log_dir)

        if os.path.exists(os.path.join(self.log_dir, "tool_logs")):
            print("tools_log_dir {} already exists".format(os.path.join(self.log_dir, "tool_logs")))
            # raise ValueError("log_dir {} already exists".format(self.log_dir))
        else:
            os.makedirs(os.path.join(self.log_dir, "tool_logs"))

        if os.path.exists(os.path.join(self.log_dir, "traces")):
            print("tools_log_dir {} already exists".format(os.path.join(self.log_dir, "traces")))
            # raise ValueError("log_dir {} already exists".format(self.log_dir))
        else:
            os.makedirs(os.path.join(self.log_dir, "traces"))

    def _initialize_task_env(self):

        work_dir = self._work_dir

        if os.path.exists(work_dir):
            for root, dirs, files in os.walk(work_dir, topdown=False):
                for name in files:
                    file_path = os.path.join(root, name)
                    if not any(subdir in file_path for subdir in ['eval_datasets', 'models', 'available_datasets']):
                        os.remove(file_path)
                for name in dirs:
                    dir_path = os.path.join(root, name)
                    if name not in ['eval_datasets', 'models', 'available_datasets']:
                        shutil.rmtree(dir_path)
        else:
            os.makedirs(work_dir)

        benchmark_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "benchmarks", self._benchmark_folder)

        if os.path.exists(os.path.join(benchmark_dir, "env")):
            for root, dirs, files in os.walk(os.path.join(benchmark_dir, "env")):
                for file in files:
                    src_path = os.path.join(root, file)
                    dst_path = os.path.join(work_dir, os.path.relpath(src_path, os.path.join(benchmark_dir, "env")))
                    os.makedirs(os.path.dirname(dst_path), exist_ok=True)
                    if not os.path.exists(dst_path):
                        shutil.copy2(src_path, dst_path)

        # find all read only files
        if os.path.exists(os.path.join(benchmark_dir, "scripts", "read_only_files.txt")):
            ignore_files = open(os.path.join(benchmark_dir, "scripts", "read_only_files.txt"), "r").read().split("\n")
            for path, subdirs, files in os.walk(os.path.join(work_dir)):

                relpath = os.path.relpath(path, work_dir)
                # filter out the files that are read only
                filenames = [os.path.join(relpath, filename) for filename in files]
                for ignore in ignore_files:
                    ignore_filenames = [n for n in filenames if fnmatch.fnmatch(n, ignore)]
                    self.read_only_files.extend(ignore_filenames)

        # init backup folder and remove all content if it exists
        if os.path.exists(os.path.join(work_dir, "backup")):
            shutil.rmtree(os.path.join(work_dir, "backup"))
        os.mkdir(os.path.join(work_dir, "backup"))

        if self.args.resume:
            shutil.rmtree(work_dir)
            resume_dir = os.path.join(self.args.resume, "env_log", "traces" , f"step_{self.args.resume_step}_files")
            print("Restoring workspace ing from {}".format(resume_dir))
            shutil.copytree(resume_dir, work_dir, symlinks=True)
            if not os.path.exists(os.path.join(work_dir, "backup")):
                os.mkdir(os.path.join(work_dir, "backup"))

    def _initialize_trace(self):
        if self.args.resume:
            print("Restoring trace from {}".format(self.args.resume))
            prev_trace = from_dict(data_class=Trace, data=json.load(open(os.path.join(self.args.resume, "env_log","trace.json"), "r")))
            print("Resetting trace to step {}".format(self.args.resume_step))
            steps = prev_trace.steps[:self.args.resume_step+1]
            t = steps[-1].timestamp
            low_level_steps = [s for s in prev_trace.low_level_steps if s.timestamp < t]
            trace = Trace(
                steps=steps,
                low_level_steps=low_level_steps,
                action_infos=self.action_infos,
                task_description=self.task_prompt,
            )
        else:   
            trace = Trace(
            steps=[],
            low_level_steps=[],
            action_infos=self.action_infos,
            task_description=self.task_prompt,
        )
        return trace
    
    def __enter__(self):
        # set time out
        def signal_handler(signum, frame):
            raise TimeoutException("Timed out!")
        signal.signal(signal.SIGALRM, signal_handler)
        signal.alarm(self.args.max_time)
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):  
        # save error message
        active = active_children()
        print(f'Active Children: {len(active)}')
        # terminate all active children
        for child in active:
            child.terminate()
        # block until all children have closed
        for child in active:
            child.join()
        # report active children
        active = active_children()
        print(f'Active Children: {len(active)}')
            
        if traceback is not None:
            print("Error message saved in error.txt")
            open(os.path.join(self.log_dir, "error.txt"), "w").write(''.join(format_exception(exc_type, exc_value, traceback)))
        open(os.path.join(self.log_dir, "overall_time.txt"), "w").write(str(time.time() - self.start_time))
        
        for sub_env in self.sub_environments.values():
            sub_env.__exit__(exc_type, exc_value, traceback)
            
    ################################# public functions ########################################

    def get_sub_environments_status(self):
        status = {}
        for attack, sub_env in self.sub_environments.items():
            status[attack] = {
                "is_final": sub_env.is_final(),
                "steps": len(sub_env.trace.steps)
            }
        return status

    def create_sub_environment(self, attack_name):
        sub_args = copy.deepcopy(self._args)
        sub_args.log_dir = os.path.join(self._log_dir, f"attackagent_{attack_name}")
        sub_args.work_dir = os.path.join(self._work_dir, f"attackagent_{attack_name}")
        sub_env = Environment(sub_args)
        
        # Copy eval_datasets, models, and available_datasets from main env to sub env
        for folder in ['eval_datasets', 'models', 'available_datasets']:
            src_folder = os.path.join(self._work_dir, folder)
            dst_folder = os.path.join(sub_env._work_dir, folder)
            if os.path.exists(src_folder):
                shutil.copytree(src_folder, dst_folder, symlinks=True, dirs_exist_ok=True)
                
        sub_env._task_prompt = self._task_prompt + f"\n\nThe above is the description for the whole task.\nYou are the AttackAgent that are initialized for {attack_name} only.\n Now perform the attack!"
        self.sub_environments[attack_name] = sub_env
        return sub_env
    
    def is_final(self):
        """Check if the task has reached a final state, either by reaching the maximum steps or time, or because the agent has submitted a final answer. """
        
        curr_step = len(self.trace.steps)
        # check if any step is final answer
        any_final_answer = any([s.action.name == "Final Answer" for s in self.trace.steps])
        return curr_step >= self.args.max_steps or any_final_answer or time.time() - self.start_time > self.args.max_time

    def execute(self, action):
        """Execute an action and return the observation."""
        
        trace = self._trace

        curr_step = len(trace.steps)
        action_name = action.name
        action_input = action.args

        if action_name == "Final Answer":
            observation = "end"

        elif self.is_final():
            observation = "The environment has shut down because the maximum number of steps or time has been reached. Please submit your final answer."

        elif action_name not in list(self.action_infos.keys()):
            actions = ", ".join(self.action_infos.keys())
            observation = f"Invalid action: {action_name}. Action did not execute. Please use one of the following actions:\n{actions}"

        else:
            # execute the action and get the observation
            log_file = os.path.join(os.path.join(self.log_dir, "tool_logs") , f"step_{curr_step}_tool_log.log")
            usage = ",\n            ".join([f"{k}: [{v}]" for k, v in self.action_infos[action_name].usage.items()])
            usage = f"""{{
            {usage}
}}"""
            invalid_action_error = f"The action input for {action_name} needs to be a valid json with proper entries. You may have missed the comma between entries. Please use the correct format and try again:\n{usage}"

            if isinstance(action_input, dict):
                try:
                    observation = self.action_infos[action_name].function(**action_input, log_file=log_file, trace=trace, **self.static_kwargs_for_tools)
                except TooLongPromptError:
                    observation="EnvError: too long input for the tool"
                except LLMError as e:
                    observation = "LLMError: " + e.message
                except EnvException as e:
                    observation = "EnvError: " + e.message
                except TypeError as e:
                    print("Step: ", curr_step, file=sys.stderr)
                    print(e, file=sys.stderr)
                    print(action_input, file=sys.stderr)
                    observation = "EnvError: " + invalid_action_error
                except TimeoutException as e:
                    raise e
                except Exception as e:
                    # should not happen
                    import traceback
                    error_traceback = traceback.format_exc()
                    print("Step: ", curr_step, file=sys.stderr)
                    print(error_traceback, file=sys.stderr)
                    observation = f"EnvError: Error executing {action_name}. Error details: {str(e)}\nTraceback:\n{error_traceback}"
            else:
                observation = invalid_action_error


        step_time = time.time()

        trace.steps.append(Step(action, observation, step_time))

        self.save(curr_step)
        return observation

    def save(self, curr_step):
        """ Save the trace and snapshot of the workspace folder """     
        with open(os.path.join(self.log_dir, f"trace.json"), "w") as f:
            json.dump(self.trace, f, indent=4, cls=EnhancedJSONEncoder)

        # Save sub-environments
        for attack, sub_env in self.sub_environments.items():
            sub_env.save(curr_step)
        ##### save a snapshot of the current step
        save_folder = os.path.join(self.log_dir, f"traces/step_{curr_step}_files")
        if os.path.exists(save_folder):
            shutil.rmtree(save_folder)
        os.makedirs(save_folder)

        # save files in the folder that are not read only
        for path, subdirs, files in os.walk(os.path.join(self.work_dir)):

            relpath = os.path.relpath(path, self.work_dir)
            dest = os.path.join(save_folder, relpath)

            for file_name in files:
                file_path = os.path.join(relpath, file_name)
                if file_path not in self.read_only_files:
                    # check wether the file to copy is part of self.log_dir
                    if  os.path.abspath(os.path.join(self.work_dir, file_path)).startswith(os.path.abspath(self.log_dir.split("/env_log")[0])):
                        continue                    
                    if not os.path.exists(dest):
                        os.makedirs(dest)            
                    shutil.copyfile(os.path.join(self.work_dir, file_path), os.path.join(save_folder, file_path))

    ############## for logging convenience ##############

    def get_task_description(self):
        return self._task_prompt, self._benchmark_folder

    @property
    def low_level_actions(self):
        return list(filter(lambda x: x.is_primitive, self.action_infos.values()))

    @property
    def high_level_actions(self):
        return list(filter(lambda x: not x.is_primitive, self.action_infos.values()))

    def print_action(self, entries):
        return "".join([ k + ": " + v for k,v in  entries.items()])