context.py

from concurrent.futures import ThreadPoolExecutor, as_completed
import json
import os.path
import subprocess
import traceback
from functools import wraps

import tqdm
from typing import Callable

from common_config import get_context_save_dir, get_D4J_ROOT, get_vul_context_save_dir, \
    get_vul_repo_root
from utils import find_all_java_files
from vul4j import get_vul_source_code_root, get_vul_test_root


def get_imports_from_java_file(file_path: str):
    encodings = ['utf-8', 'latin-1', 'cp1252', 'gbk']
    imports = []

    for encoding in encodings:
        try:
            with open(file_path, 'r', encoding=encoding) as f:
                for line in f:
                    if 'import ' in line:
                        imports.append(line.strip())
            print(f"Success with encoding: {encoding}")
            return imports
        except UnicodeDecodeError:
            continue

    raise ValueError(f"cannot use encoding: {encodings}")


def analyze_java_file_and_save(bug_project, bug_id, file_path):
    if not file_path.endswith('.java'):
        print('file_path must endswith .java')
        return

    save_dir = f'{get_context_save_dir()}/{bug_project}_{bug_id}'
    save_file = file_path.replace(get_D4J_ROOT() + '/', '').replace('/', '_').replace(".java", ".json")
    os.makedirs(save_dir, exist_ok=True)

    if os.path.exists(f'{save_dir}/{save_file}'):
        print('file exists')
        return

    current_dir = os.path.dirname(os.path.abspath(__file__))
    result = subprocess.run(
        ['java', '-jar', f'{current_dir}/artifacts/FileParser-1.0-SNAPSHOT-jar-with-dependencies.jar', file_path],
        capture_output=True,
        text=True,
        timeout=30
    )
    json_result = json.loads(result.stdout)
    json_result['imports'] = get_imports_from_java_file(file_path)

    with open(f'{save_dir}/{save_file}', 'w') as f:
        f.write(json.dumps(json_result, indent=4))


def analyze_java_file_and_save_for_vul(vul_id, file_path):
    if not file_path.endswith('.java'):
        print('file_path must endswith .java')
        return

    save_dir = f'{get_vul_context_save_dir()}/{vul_id}'
    save_file = file_path.replace(get_vul_repo_root() + '/', '').replace('/', '_').replace(".java", ".json")
    os.makedirs(save_dir, exist_ok=True)

    if os.path.exists(f'{save_dir}/{save_file}'):
        print('file exists')
        return

    current_dir = os.path.dirname(os.path.abspath(__file__))
    result = subprocess.run(
        ['java', '-jar', f'{current_dir}/artifacts/FileParser-1.0-SNAPSHOT-jar-with-dependencies.jar', file_path],
        capture_output=True,
        text=True,
        timeout=30
    )

    json_result = json.loads(result.stdout)
    json_result['imports'] = get_imports_from_java_file(file_path)

    with open(f'{save_dir}/{save_file}', 'w') as f:
        f.write(json.dumps(json_result, indent=4))


class_skeleton_template = """
======================== Begin Class ========================
Full Qualified Name: {full_qualified_name}

File Path: {file_path}

Imports: 
{imports}

Fields: 
{fields}

Methods:
{methods}

======================== End Class ========================
"""


def tool(func: Callable) -> Callable:
    func.is_tool = True

    @wraps(func)
    def wrapper(*args, **kwargs):
        return func(*args, **kwargs)

    return wrapper


def process_single_java_file(bug_project, bug_id, java_file):
    try:
        analyze_java_file_and_save(bug_project, bug_id, java_file)
    except Exception as e:
        print("👇 Full Traceback:")
        traceback.print_exc()
        print("-" * 60)


def process_single_java_file_for_vul(vul_id, java_file):
    try:
        analyze_java_file_and_save_for_vul(vul_id, java_file)
    except Exception as e:
        print(f'mistake happens parsing {java_file}')
        print("👇 Full Traceback:")
        traceback.print_exc()
        print("-" * 60)


class ContextTool:
    def __init__(self,
                 bug_project,
                 bug_id,
                 need_init=True,
                 is_vul=False,
                 vul_id=''):
        if not is_vul:
            bug_project_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/source/org'
            if not os.path.exists(bug_project_root_dir):
                bug_project_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/src/org'
            if not os.path.exists(bug_project_root_dir):
                bug_project_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/src/java/org'  # cli
            if not os.path.exists(bug_project_root_dir):
                bug_project_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/src/main/java/org'  # cli_30
            if not os.path.exists(bug_project_root_dir):
                bug_project_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/src/main/java/com'  # jackson
            if not os.path.exists(bug_project_root_dir):
                bug_project_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/src/com'  # closure
            if not os.path.exists(bug_project_root_dir) and bug_project == 'Gson':
                bug_project_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/gson/src/main/java/com'  # gson
            if not os.path.exists(bug_project_root_dir):
                raise Exception(f'cannot find root dir of the buggy project {bug_project}')
        else:
            bug_project_root_dir = get_vul_repo_root() + '/' + vul_id + '/' + get_vul_source_code_root(vul_id)
            print(bug_project_root_dir)

        if not is_vul:
            test_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/tests/org'
            if not os.path.exists(test_root_dir):
                test_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/test/org'
            if not os.path.exists(test_root_dir):
                test_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/src/tests/org'
            if not os.path.exists(test_root_dir):
                test_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/src/test/org'
            if not os.path.exists(test_root_dir):
                test_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/src/tests/java/org'
            if not os.path.exists(test_root_dir):
                test_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/src/test/java/org'
            if not os.path.exists(test_root_dir):
                test_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/src/tests/java/com'
            if not os.path.exists(test_root_dir):
                test_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/src/test/java/com'
            if not os.path.exists(test_root_dir):
                test_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/tests/com'
            if not os.path.exists(test_root_dir):
                test_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/test/com'
            if not os.path.exists(test_root_dir) and bug_project == 'Gson':
                test_root_dir = f'{get_D4J_ROOT()}/{bug_project}_{bug_id}/gson/src/test/java/com'  # gson
            if not os.path.exists(test_root_dir):
                raise Exception(f'cannot find test dir of the buggy project {bug_project}')
        else:
            test_root_dir = get_vul_repo_root() + '/' + vul_id + '/' + get_vul_test_root(vul_id)
            print(test_root_dir)

        self.test_root_dir = test_root_dir
        self.bug_project_root_dir = bug_project_root_dir
        self.bug_project = bug_project
        self.bug_id = bug_id
        self.vul_id = vul_id
        if not is_vul:
            self.context_dir = f'{get_context_save_dir()}/{bug_project}_{bug_id}'
        else:
            self.context_dir = f'{get_vul_context_save_dir()}/{vul_id}'

        if need_init:
            source_java_files = find_all_java_files(bug_project_root_dir)
            if self.test_root_dir is not None:
                test_java_files = find_all_java_files(self.test_root_dir)
            else:
                test_java_files = []

            all_tasks = []

            for java_file in source_java_files:
                if is_vul:
                    all_tasks.append((self.vul_id, java_file))
                else:
                    all_tasks.append((self.bug_project, self.bug_id, java_file))

            for java_file in test_java_files:
                if is_vul:
                    all_tasks.append((self.vul_id, java_file))
                else:
                    all_tasks.append((self.bug_project, self.bug_id, java_file))

            with ThreadPoolExecutor(max_workers=16) as executor:

                if is_vul:
                    futures = [
                        executor.submit(process_single_java_file_for_vul, vul_id, java_file)
                        for vul_id, java_file in all_tasks
                    ]
                else:
                    futures = [
                        executor.submit(process_single_java_file, bug_project, bug_id, java_file)
                        for bug_project, bug_id, java_file in all_tasks
                    ]

                for future in tqdm.tqdm(as_completed(futures),
                                        total=len(futures), desc='Processing Java files'):
                    pass

    def find_class_info_by_full_qualified_name(self, class_full_qualified_name):
        class_info = None
        for x in os.listdir(self.context_dir):
            path = self.context_dir + os.path.sep + x
            if class_full_qualified_name.replace('.', '_') + '.json' in x:
                with open(path, 'r') as f:
                    json_obj = json.load(f)
                if json_obj['fullQualifiedName'] == class_full_qualified_name:
                    # find it
                    class_info = json_obj
                    break
        return class_info

    @tool
    def get_project_structure(self, file_included=None):
        proj_structure = f"Source Code:\n{self._get_project_structure(self.bug_project_root_dir)}"
        if self.test_root_dir is not None:
            proj_structure += f'\n\nTest Code:\n{self._get_project_structure(self.test_root_dir)}'
        return proj_structure

    def _get_project_structure(self, target_dir, file_included=None):
        from io import StringIO
        import sys
        old_stdout = sys.stdout
        sys.stdout = captured_output = StringIO()

        if file_included is None:
            file_included = ['.java']

        file_included = [ext.lower() for ext in file_included]

        excluded_dirs = {
            'build', 'target', 'bin', '.git', '.idea', '__pycache__',
            'node_modules', 'dist', 'out'
        }

        def _tree(dir_path: str, prefix: str = '', is_last: bool = True) -> None:
            """
            print tree recursively
            """
            if not os.path.exists(dir_path):
                print(f"{prefix}└── [Error: Path does not exist]")
                return

            try:
                entries = sorted(os.listdir(dir_path))
            except PermissionError:
                print(f"{prefix}└── [Permission Denied]")
                return

            dirs = []
            files = []
            for e in entries:
                full_path = os.path.join(dir_path, e)
                if os.path.isdir(full_path):
                    if e not in excluded_dirs:
                        dirs.append(e)
                else:
                    ext = os.path.splitext(e)[1].lower()
                    if ext in file_included:
                        files.append(e)

            items = dirs + files
            if not items:
                return

            pointers = [('├── ', False)] * (len(items) - 1) + [('└── ', True)]

            for ptr, is_last_item in pointers:
                item = items.pop(0)
                full_path = os.path.join(dir_path, item)
                print(f"{prefix}{ptr}{item}")

                if item in dirs:
                    extension = '│   ' if not is_last_item else '    '
                    _tree(full_path, prefix + extension, is_last=True)

        print(f"{os.path.basename(target_dir)}/")
        _tree(target_dir, prefix='', is_last=True)

        sys.stdout = old_stdout
        return captured_output.getvalue()

    @tool
    def get_class_skeleton(self, class_full_qualified_name):

        class_info = self.find_class_info_by_full_qualified_name(class_full_qualified_name)
        # class not found
        if class_info is None:
            return None

        import_str = ''.join(['  ' + x + '\n' for x in class_info['imports']])
        field_str = ''.join(['  ' + f'name: {x["name"]}, type: {x["type"]}, modifiers: {x["modifiers"]}' + '\n' for x in
                             class_info['fieldList']])
        method_str = ''.join(['  ' + x['signature'] + '\n' for x in class_info['methodList']])
        return class_skeleton_template.format(
            full_qualified_name=class_info['fullQualifiedName'],
            file_path=class_info['path'],
            imports=import_str,
            fields=field_str,
            methods=method_str
        )

    @tool
    def get_field_type(self, class_full_qualified_name, field_name):
        class_info = self.find_class_info_by_full_qualified_name(class_full_qualified_name)
        # class not found
        if class_info is None:
            return None

        for x in class_info['fieldList']:
            if x['name'] == field_name:
                return x['type']
        # field not found
        return None

    @tool
    def get_local_variable_type(self, class_full_qualified_name, function_name, var_name):
        class_info = self.find_class_info_by_full_qualified_name(class_full_qualified_name)
        # class not found
        if class_info is None:
            return None

        target_method_info = None
        for x in class_info['methodList']:
            if x['name'] == function_name:
                target_method_info = x
                break

        # method not found
        if target_method_info is None:
            return None

        for y in target_method_info['localVars']:
            if y['name'] == var_name:
                return y['type']
        # variable not found
        return None

    @tool
    def get_function_body(self, class_full_qualified_name, function_name):
        class_info = self.find_class_info_by_full_qualified_name(class_full_qualified_name)
        # class not found
        if class_info is None:
            return None
        result = ''
        for x in class_info['methodList']:
            if x['name'] == function_name:
                result += x['signature']
                result += x['content']
                result += '\n\n'

        # method not found
        return result


if __name__ == '__main__':
    context_tool = ContextTool(bug_project='Chart',
                               bug_id=12)
    print(context_tool.get_project_structure())