Sage will provide the foundation for scanning and processing Ansible projects as a code.
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

sage_scan/__init__.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

__version__ = "0.0.1"

sage_scan/__version__.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

sage_scan/custom_scan/concurrency/__init__.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from sage_scan.pipeline import SagePipeline
import os
import argparse
import json
import redis
from sage_scan.tools.src_rebuilder import write_result, prepare_source_dir


REDIS_SERVER_URL = os.getenv("REDIS_SERVER_URL", 'localhost')
REDIS_SERVER_PORT = int(os.getenv("REDIS_SERVER_PORT", '6379'))
REDIS_SERVER_DB = int(os.getenv("REDIS_SERVER_DB", '0'))


def check_if_result_exists(result_dir: str, src_type: str, repo_name: str):
    objects_path = os.path.join(result_dir, src_type, repo_name, "sage-objects.json")
    exists = os.path.exists(objects_path)
    return exists


# TODO: imeplement this
def is_task_queued(redis_client: redis.Redis, queue_name: str, src_type: str, repo_name: str):
    # redis_client.lrange(queue_name, 0, -1)
    return False


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="TODO")
    parser.add_argument("-f", "--filepath", default="", help="filepath to custom scan input JSON")
    

    parser = argparse.ArgumentParser(description="TODO")
    parser.add_argument("-t", "--source-type", help='source type (e.g."GitHub-RHIBM")')
    parser.add_argument("-s", "--source-json", help='source json file path (e.g. "/tmp/RH_IBM_FT_data_GH_api.json")')
    parser.add_argument("-o", "--out-dir", default="./sage_concurrency_work_dir", help="output directory")
    parser.add_argument("--timeout", help="timeout seconds")
    parser.add_argument("--resume", action="store_true", help="if true, only tasks without existing results are registered (queued tasks are also skipped)")
    parser.add_argument("--dry-run", action="store_true", help="if true, it does not register tasks actually (normally use this mode with `--save` option)")
    parser.add_argument("--save", action="store_true", help="if true, save a json file which is a list of registered projects")
    args = parser.parse_args()

    work_dir = args.out_dir
    in_container_work_dir = "/work_dir"
    src_type = args.source_type
    src_json = args.source_json
    src_rb_dir = os.path.join(work_dir, "src_rb")
    path_list_dir = os.path.join(work_dir, "path_list")
    result_dir = os.path.join(work_dir, "results")

    os.makedirs(work_dir, exist_ok=True)
    os.makedirs(src_rb_dir, exist_ok=True)
    os.makedirs(path_list_dir, exist_ok=True)
    os.makedirs(result_dir, exist_ok=True)

    adir = os.path.join(src_rb_dir, src_type)
    if not os.path.exists(adir) or len(os.listdir(adir)) == 0:
        outfile = os.path.join(path_list_dir, f"path-list-{src_type}.txt")
        path_list = prepare_source_dir(adir, src_json)
        write_result(outfile, path_list)

    repo_names = set()
    if src_json:
        with open(src_json, "r") as f:
            records = f.readlines()
        
        for record in records:
            r = json.loads(record)
            if "repo_name" in r:
                repo_names.add(r.get("repo_name"))
            elif "namespace_name" in r:
                repo_names.add(r.get("namespace_name"))

    redis_client = redis.Redis(
        host=REDIS_SERVER_URL,
        port=REDIS_SERVER_PORT,
        db=REDIS_SERVER_DB,
    )

    count = 0
    registerd_lines = []
    for repo_name in repo_names:
        if args.resume:
            if check_if_result_exists(result_dir, src_type, repo_name):
                # skip because the result exists
                continue
            # print(repo_name)
            
            # TODO: imeplement this
            # if is_task_queued(redis_client, "queue", src_type, repo_name):
            #     # skip because the task is found in the current queue
            #     continue

        target_dir = os.path.join(in_container_work_dir, "src_rb", src_type, repo_name)
        
        task_input = {
            "dir_path": target_dir,
            "source_type": src_type,
            "repo_name": repo_name,
        }
        if args.timeout:
            task_input["timeout"] = int(args.timeout)

        task_input_str = json.dumps(task_input)
        if not args.dry_run:
            redis_client.rpush("queue", task_input_str)
        count += 1
        registerd_lines.append(task_input_str + "\n")

    dry_run_msg = ""
    if args.dry_run:
        dry_run_msg = " [DRY_RUN]"

    print(f"scan for {count} projects has been registered.{dry_run_msg}")

    if args.save:
        registerd_file = "registered_projects.json"
        with open(registerd_file, "w") as file:
            file.write("".join(registerd_lines))

sage_scan/custom_scan/concurrency/register.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from sage_scan.pipeline import SagePipeline
import os
import time
import traceback
from celery import Celery



dp = SagePipeline()


REDIS_SERVER_URL = os.getenv("REDIS_SERVER_URL", 'localhost')
REDIS_SERVER_PORT = int(os.getenv("REDIS_SERVER_PORT", '6379'))
REDIS_SERVER_DB = int(os.getenv("REDIS_SERVER_DB", '0'))

SINGLE_SCAN_TIMEOUT = int(os.getenv("SINGLE_SCAN_TIMEOUT", '450'))

CELERY_BROKER_URL = f"redis://{REDIS_SERVER_URL}:{REDIS_SERVER_PORT}/1"
CELERY_RESULT_BACKEND = f"redis://{REDIS_SERVER_URL}:{REDIS_SERVER_PORT}/2"


celery = Celery(__name__)
celery.conf.broker_url = CELERY_BROKER_URL
celery.conf.result_backend = CELERY_RESULT_BACKEND


@celery.task(name='tasks.scan')
def scan(dir_path: str, source_type: str, repo_name: str, out_dir: str) -> dict:
    source = {
        "type": source_type,
        "repo_name": repo_name,
    }
    result = {
        "success": None,
        "error": None,
        "begin": time.time(),
        "end": None,
    }
    err = None
    try:
        _ = dp.run(
            target_dir=dir_path,
            output_dir=out_dir,
            source=source,
        )
    except Exception:
        err = traceback.format_exc()
    finally:
        success = False if err else True
        result["success"] = success
        if not success:
            result["error"] = err
        result["end"] = time.time()
    return result

sage_scan/custom_scan/concurrency/tasks.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import argparse
import json
import redis
import time

from tasks import scan


REDIS_SERVER_URL = os.getenv("REDIS_SERVER_URL", 'localhost')
REDIS_SERVER_PORT = int(os.getenv("REDIS_SERVER_PORT", '6379'))
REDIS_SERVER_DB = int(os.getenv("REDIS_SERVER_DB", '0'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="TODO")
    parser.add_argument("-o", "--out-dir", default="/work_dir/results", help="output directory for the rule evaluation result")
    args = parser.parse_args()

    os.environ['SAGE_CONTENT_ANALYSIS_OUT_DIR'] = args.out_dir

    redis_client = redis.Redis(
        host=REDIS_SERVER_URL,
        port=REDIS_SERVER_PORT,
        db=REDIS_SERVER_DB,
    )

    while True:
        task_input_str = redis_client.lpop('queue')
        if not task_input_str:
            time.sleep(3)
            continue
        task_input = json.loads(task_input_str)
        dir_path = task_input["dir_path"]
        source_type = task_input["source_type"]
        repo_name = task_input["repo_name"]
        out_dir = os.path.join(args.out_dir, source_type, repo_name)
        timeout = task_input.get("timeout", None)
        
        kwargs = {
            "dir_path": dir_path,
            "source_type": source_type,
            "repo_name": repo_name,
            "out_dir": out_dir,
        }

        if timeout:
            result = scan.apply_async(kwargs=kwargs, time_limit=timeout)
        else:
            result = scan.apply_async(kwargs=kwargs)
sage_scan/custom_scan/concurrency/trigger.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from sage_scan.pipeline import SagePipeline
import os
import argparse


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="TODO")
    parser.add_argument("-d", "--dir", help='root direcotry for scan')
    parser.add_argument("-t", "--source-type", help='source type (e.g."GitHub-RHIBM")')
    parser.add_argument("-r", "--repo-name", help='repo name (e.g."IBM/Ansible-OpsnShift-Provisioning")')
    parser.add_argument("-o", "--out-dir", default="", help="output directory for the rule evaluation result")
    args = parser.parse_args()

    source = {}
    if args.source_type:
        source["type"] = args.source_type
    if args.repo_name:
        source["repo_name"] = args.repo_name

    dp = SagePipeline()
    os.environ['SAGE_CONTENT_ANALYSIS_OUT_DIR'] = args.out_dir
    result = dp.run(
        target_dir=args.dir,
        output_dir=args.out_dir,
        source=source,
    )

sage_scan/custom_scan/custom_scan.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from sage_scan.pipeline import SagePipeline
import os
import json
import time
import argparse
from sage_scan.tools.src_rebuilder import write_result, prepare_source_dir
from ansible_risk_insight.scanner import ARIScanner, Config


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="TODO")
    parser.add_argument("-t", "--source-type", help='source type (e.g."GitHub-RHIBM")')
    parser.add_argument("-s", "--source-json", help='source json file path (e.g. "/tmp/RH_IBM_FT_data_GH_api.json")')
    parser.add_argument("-o", "--out-dir", help="output directory")
    parser.add_argument("-p", "--project-list", help="project list")
    parser.add_argument("--yml-inventory-only", action="store_true", help="yml inventory only mode")
    args = parser.parse_args()

    work_dir = args.out_dir
    src_type = args.source_type
    src_json = args.source_json
    project_list = args.project_list
    yml_inventory_mode = args.yml_inventory_only
    src_rb_dir = os.path.join(work_dir, "src_rb")
    path_list_dir = os.path.join(work_dir, "path_list")
    result_dir = os.path.join(work_dir, "results")

    os.makedirs(work_dir, exist_ok=True)
    os.makedirs(src_rb_dir, exist_ok=True)
    os.makedirs(path_list_dir, exist_ok=True)
    os.makedirs(result_dir, exist_ok=True)

    adir = os.path.join(src_rb_dir, src_type)
    if not os.path.exists(adir) or len(os.listdir(adir)) == 0:
        outfile = os.path.join(path_list_dir, f"path-list-{src_type}.txt")
        path_list = prepare_source_dir(adir, src_json)
        write_result(outfile, path_list)

    if src_json:
        with open(src_json, "r") as f:
            records = f.readlines()
        repo_names = set()
        for record in records:
            r = json.loads(record)
            if "repo_name" in r:
                repo_names.add(r.get("repo_name"))
            if "namespace_name" in r:
                repo_names.add(r.get("namespace_name"))

    if project_list:
        with open(project_list, "r") as f:
            repo_names = [s.rstrip() for s in f.readlines()]

    total = len(repo_names)
    count = 0

    out_scope = [
        "IBM/playbook-integrity-operator",
        "RedHatOfficial/ansible-role-rhv4-rhvh-stig",
        "confluent.platform",
        "bosh-io/releases-index"
    ]

    dp = SagePipeline()

    timer_path = "/tmp/custom-scan-all-timer.json"
    for repo_name in repo_names:
        if repo_name in out_scope:
            print(f"skip {repo_name} ({count}/{total})")
            count += 1
            continue

        start = time.time()
        tdir = os.path.join(src_rb_dir, src_type, repo_name)
        odir = os.path.join(result_dir, src_type, repo_name)
        if os.path.exists(os.path.join(odir, "ftdata.json")):
            count += 1
            continue

        # why needed?
        os.environ["SAGE_CONTENT_ANALYSIS_OUT_DIR"] = odir

        print(f"scanning {repo_name} ({count}/{total})")

        dp.run(
            target_dir=tdir,
            output_dir=odir,
            source={"type": src_type, "repo_name": repo_name},
            yml_inventory_only=yml_inventory_mode,
        )
        count += 1

        end = time.time()
        elapsed = end - start
        timer_record = {
            "repo_name": repo_name,
            "elapsed": elapsed,
        }
        with open(timer_path, "a+") as file:
            file.write(json.dumps(timer_record) + "\n")
        
sage_scan/custom_scan/custom_scan_all.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os

from sage_scan.pipeline import SagePipeline


EXAMPLE = '''
---
- name: AWS EC2 Cloud Operations
  hosts: localhost
  gather_facts: false

  vars:
    image_id: ami-016eb5d644c333ccb
    key_name: app-keypair
    security_groups:
      - ansibull-sg

  tasks:
    - name: Create a key pair in us-east-2 called app-keypair
      ec2_key:
        name: app-keypair
        region: us-east-2
      register: ec2_key

    - name: Create a t3.medium instance called app-instance-01
      ec2_instance:
        key_name: "{{ key_name }}"
        name: app-instance-01
        image_id: "{{ image_id }}"
        instance_type: t3.medium
        vpc_subnet_id: "{{ vpc_subnet_id }}"
        security_group: "{{ security_group }}"
        network:
          assign_public_ip: true
        tags:
          Environment: Testing
      register: ec2
'''


def main():
    input_yaml = EXAMPLE

    sp = SagePipeline()

    # process the input playbook by SagePipeline
    project = sp.run(
        raw_yaml=input_yaml,
    )

    # Here you can do anything with `project` which is a SageProject object (=pipeline output)


if __name__ == '__main__':
    main()


sage_scan/custom_scan/custom_scan_str.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import jsonpickle
import threading
from datetime import date, time
from dataclasses import dataclass, field
import ruamel.yaml

from ansible_risk_insight.models import (
    AnsibleRunContext,
    RunTargetType,
    Rule,
    RuleResult,
    Severity,
    Collection,
    Role,
    Repository,
    Playbook,
    TaskFile,
    Task,
    TaskCall,
    Variable,
    VariableType,
)
import ansible_risk_insight.yaml as ariyaml
import os
import re

import yaml
try:
    # if `libyaml` is available, use C based loader for performance
    import _yaml  # noqa: F401
    from yaml import CSafeDumper as Dumper
except Exception:
    # otherwise, use Python based loader
    from yaml import SafeDumper as Dumper

# Copied from distutils.util.strtobool
def strtobool(val):
    """Convert a string representation of truth to true (1) or false (0).
    True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values
    are 'n', 'no', 'f', 'false', 'off', and '0'.  Raises ValueError if
    'val' is anything else.
    """
    val = val.lower()
    if val in ('y', 'yes', 't', 'true', 'on', '1'):
        return 1
    elif val in ('n', 'no', 'f', 'false', 'off', '0'):
        return 0
    else:
        raise ValueError("invalid truth value %r" % (val,))

# Set default representer
def default_representer(dumper, data):
    return dumper.represent_scalar('tag:yaml.org,2002:str', serialize(data))
yaml.representer.SafeRepresenter.add_representer(None, default_representer)


thread_id = threading.get_native_id()
in_parallel = strtobool(os.getenv("SAGE_CONTENT_ANALYSIS_PARALLEL", "False"))
do_save_scan_result = strtobool(os.getenv("SAGE_SAVE_SCAN_RESULT", "False"))
out_dir = os.getenv("SAGE_CONTENT_ANALYSIS_OUT_DIR", "/tmp/ftdata")
file_suffix = f"_{thread_id}" if in_parallel else ""
default_task_context_data_filepath = os.path.join(out_dir, f"task_context_data{file_suffix}.json")
default_scan_res_filepath = os.path.join(out_dir, f"scan_result{file_suffix}.json")

vars_json_length_threshold = 1000
max_context_node_size = 50
max_yaml_str_size = 3000


def find_repo_info(ctx: AnsibleRunContext):
    ctx_1 = ctx.copy()  # --> reset current

    scan_data = ctx_1.scan_metadata
    scan_type = scan_data.get("type")
    scan_path = scan_data.get("name")

    src_path_list_file = os.getenv('FTDATA_SRC_PATH_LIST')

    with open(src_path_list_file, "r") as file:
        lines = file.readlines()

    input_list = []
    for line in lines:
        data = json.loads(line)
        input_list.append(data)

    repo_name = ""
    repo_type = ""
    license = ""
    source = ""
    path = ""

    for input in input_list:
        input_type = input.get("type")
        if scan_type != input_type:
            continue

        input_path = ""
        if scan_type == "role":
            input_path = input.get("role_path")
        else:
            input_path = input.get("target_file")
        if scan_path.endswith(input_path):
            repo_name = input.get("_repo_name")
            repo_type = input.get("_repo_type")
            source = input.get("_source")
            license = input.get("_license")
            path = input.get("_path")

    return repo_type, repo_name, license, source, path


def make_yaml_before_task(ctx: AnsibleRunContext, task: Task) -> list:
    ctx_1 = ctx.copy()  # --> reset current
    yaml_before_task = ""
    filepath_of_this_task = task.defined_in
    for node in ctx_1:
        if not isinstance(node.spec, Playbook) and not isinstance(node.spec, TaskFile):
            continue
        if node.spec.defined_in != filepath_of_this_task:
            continue

        file_obj = node.spec
        if not hasattr(file_obj, "yaml_lines"):
            # if the file does not have `yaml_lines`, give up here
            break

        if not task.line_num_in_file or len(task.line_num_in_file) != 2:
            # if the task does not have `line_num_in_file`, give up here
            break

        # the num of line where the task starts
        task_line_num = task.line_num_in_file[0]

        lines = file_obj.yaml_lines.splitlines()
        if len(lines) < task_line_num:
            # if file lines are less than the task_line_num, give up here
            break

        yaml_before_task = "\n".join(lines[:task_line_num-1])
        break

    # we omit yaml_before_task if it is huge so that we can avoid data size issue
    # if len(yaml_before_task) > max_yaml_str_size:
    #     yaml_before_task = ""
    return yaml_before_task


def get_used_modules(ctx, task):
    ctx_1 = ctx.copy()  # --> reset current
    used_modules = []
    for node in ctx_1:
        # skip non-task node
        if node.type != RunTargetType.Task:
            continue
        # if node is the target task, exit loop
        if node.spec.key == task.spec.key:
            break

        module = get_module_name(node)
         # if the module is empty, skip it
        if not module:
            continue
        # if the module is ansible.builtin one, skip it
        if module.startswith("ansible.builtin."):
            continue
        # skip assert (it is covered by ansible.builtin condition though)
        if module.endswith("assert"):
            continue
        # if the module is already in used_modules, skip it
        if module in used_modules:
            continue
        used_modules.append(module)
    return used_modules


def get_name_sequence(ctx, task):
    ctx_1 = ctx.copy()  # --> reset current
    task_node_id = task.node_id
    def get_parent_node_id(node_id):
        if "." in node_id:
            return node_id.rsplit(".", 1)[0]
        # if node_id is like "3", then we regard "" (empty string) as the parent id
        return ""
    task_parent_node_id = get_parent_node_id(task_node_id)
    task_node_depth = task.depth
    name_seq = []
    for node in ctx_1:
        skip = True
        if node.depth < task_node_depth:
            if task_node_id.startswith(node.node_id):
                skip = False
        elif node.depth == task_node_depth:
            if get_parent_node_id(node.node_id) == task_parent_node_id:
                skip = False

        if skip:
            continue

        name = ""
        if node.type == RunTargetType.Role:
            if isinstance(node.spec.metadata, dict):
                name = node.spec.metadata.get("galaxy_info", {}).get("description", "")
        elif node.type == RunTargetType.Play:
            name = node.spec.name
        elif node.type == RunTargetType.Task:
            name = node.spec.name
        if node.spec.key == task.spec.key:
            break
        if not name:
            continue
        name_seq.append(name)

    # we truncate older context if the name sequence is too large
    if len(name_seq) > max_context_node_size:
        name_seq = name_seq[-max_context_node_size:]

    return name_seq

def get_parents(ctx, task):
    parents = {
        "role": [],
        "taskfile": [],
    }

    ctx_1 = ctx.copy()  # --> reset current
    task_node_id = task.node_id
    def get_parent_node_id(node_id):
        if "." in node_id:
            return node_id.rsplit(".", 1)[0]
        # if node_id is like "3", then we regard "" (empty string) as the parent id
        return ""
    task_parent_node_id = get_parent_node_id(task_node_id)
    task_node_depth = task.depth

    for node in ctx_1:
        skip = True
        if node.depth < task_node_depth:
            if task_node_id.startswith(node.node_id):
                skip = False
        elif node.depth == task_node_depth:
            if get_parent_node_id(node.node_id) == task_parent_node_id:
                skip = False

        if skip:
            continue

        if node.type == RunTargetType.Role:
            pr = {}
            pr["name"] = node.spec.name
            if node.spec.metadata:
                pr["description"] = node.spec.metadata.get("galaxy_info", {}).get("description", "")
            parents["role"].append(pr)
        elif node.type == RunTargetType.TaskFile:
            # if current_file != node.spec.defined_in:
            parents["taskfile"].append(node.spec.defined_in)

    return parents


def get_used_vars(ctx, task):
    ctx_1 = ctx.copy()  # --> reset current

    # we don't use `task.variable_use` here because it contas vars used by the target task
    # instead, we use `variable_use` of the previous task if the one exist
    previous_task = None
    for node in ctx_1:
        if node.type != RunTargetType.Task:
            continue
        if node.spec.key == task.spec.key:
            break
        previous_task = node

    # if there is no task before the target, return no vars
    if not previous_task:
        return [], False

    used_var_names = list(previous_task.variable_use.keys())
    omitted = False
    used_var_names_str = json.dumps(used_var_names)
    if len(used_var_names_str) > vars_json_length_threshold:
        used_var_names = []
        omitted = True

    return used_var_names, omitted


def get_defined_vars(task):
    defined_var_names = list(task.variable_set.keys())
    omitted = False
    defined_var_names_str = json.dumps(defined_var_names)
    if len(defined_var_names_str) > vars_json_length_threshold:
        defined_var_names = []
        omitted = True

    return defined_var_names, omitted


def get_defined_vars_in_parent(ctx, task):
    ctx_1 = ctx.copy()  # --> reset current
    defined_vars_name_value = {}
    used_vars_name_value = {}
    defined_var_names = []

    registered_var = VariableType.RegisteredVars.name

    for node in ctx_1:
        # skip non-task node
        if node.type != RunTargetType.Task:
            continue
        # if node is the target task, exit loop
        if node.spec.key == task.spec.key:
            break

        defined_var_names = list(node.variable_set.keys())

        for d_vars in node.variable_set.values():
            if type(d_vars) == list:
                for dv in d_vars:
                    if dv.value is not None and dv.type.name != registered_var:
                        defined_vars_name_value[dv.name] = dv.value
            elif d_vars:
                if d_vars.value is not None and dv.type.name != registered_var:
                    defined_vars_name_value[d_vars.name] = d_vars.value

        for u_vars in node.variable_use.values():
            if type(u_vars) == list:
                for uv in u_vars:
                    if uv.value is not None and uv.type.name != registered_var:
                        used_vars_name_value[uv.name] = uv.value
            elif u_vars:
                if u_vars.value is not None and uv.type.name != registered_var:
                    used_vars_name_value[u_vars.name] = u_vars.value

    # get variable values for defined variable names
    found_var_name_values = {}
    for dvn in defined_var_names:
        if dvn in defined_vars_name_value:
            found_var_name_values[dvn] = defined_vars_name_value[dvn]
        elif dvn in used_vars_name_value:
            found_var_name_values[dvn] = used_vars_name_value[dvn]
        else:
            found_var_name_values[dvn] = ""
    return found_var_name_values

def get_used_vars_in_target_task(task):
    var_names = []
    vars = task.args.vars
    for var in vars:
        var_names.append(var.name)
    return var_names

def get_dependency(ctx):
    parent = ctx.parent
    coll_dependencies = []
    role_dependencies = []
    if isinstance(parent, Collection):
        # there was a bug around "collection.dependency"
        # so we get dep info from metadata instead
        # dependency = parent.dependency
        original = parent.metadata.get("collection_info", {}).get("dependencies", {})
        if original:
            coll_dependencies = list(original.keys())
    elif isinstance(parent, Role):
        original = parent.dependency
        if original and isinstance(original, dict):
            coll_dependencies = original.get("collections", [])
            role_dependencies = original.get("roles", [])
    return coll_dependencies, role_dependencies


def get_module_name(task):
    module = task.get_annotation("module.correct_fqcn", "")

    if not module and task.spec.module:
        short_name = task.spec.module.split(".")[-1]
        if short_name == "include_role":
            module = "ansible.builtin.include_role"
        elif short_name == "import_role":
            module = "ansible.builtin.import_role"
        elif short_name == "include_tasks":
            module = "ansible.builtin.include_tasks"
        elif short_name == "import_tasks":
            module = "ansible.builtin.import_tasks"
        elif short_name == "include":
            module = "ansible.builtin.include"

    if not module or "." in task.spec.module:
        module = task.spec.module

    return module


def get_metrics(ctx, task, module, used_modules, used_vars, defined_vars, coll_deps, role_deps):
    d = {
        "module_name": "",
        "module_parent": "",
        "vars_used_in_this_task": [],
        "num_of_vars_used_in_this_task": 0,
        "exact_match_in_used_modules": 0,
        "parent_match_in_used_modules": 0,
        "parent_found_in_dependencies": 0,
        "num_of_vars_found_in_used_vars": 0,
        "num_of_vars_found_in_defined_vars": 0,
        "num_of_vars_covered_by_context": 0,
    }
    d["module_name"] = module

    if "." not in module:
        return d

    parts = module.split(".")
    module_parent = ""
    if len(parts) == 2:
        module_parent = parts[0]
    elif len(parts) >= 3:
        module_parent = ".".join(parts[:2])
    d["module_parent"] = module_parent

    if module_parent == "ansible.builtin":
        return d

    exact_match = [um for um in used_modules if um == module]
    parent_match = [um for um in used_modules if um.startswith(f"{module_parent}.")]
    d["exact_match_in_used_modules"] = len(exact_match)
    d["parent_match_in_used_modules"] = len(parent_match)

    dep_count = 0
    if module_parent in coll_deps:
        dep_count = 1
    elif module_parent in role_deps:
        dep_count = 1

    d["parent_found_in_dependencies"] = dep_count

    if not task.args:
        return d

    vars_used_in_this_task = list(set([v.name.split(".")[0] for v in task.args.vars]))
    d["vars_used_in_this_task"] = vars_used_in_this_task
    num_of_vars_used_in_this_task = len(vars_used_in_this_task)
    d["num_of_vars_used_in_this_task"] = num_of_vars_used_in_this_task

    if num_of_vars_used_in_this_task == 0:
        return d

    found_used_var_count = 0
    found_defined_var_count = 0
    covered_var_count = 0
    for var_name in vars_used_in_this_task:
        covered = False
        if var_name in used_vars:
            found_used_var_count += 1
            covered = True

        if var_name in defined_vars:
            found_defined_var_count += 1
            covered = True

        if covered:
            covered_var_count += 1

    d["num_of_vars_found_in_used_vars"] = found_used_var_count
    d["num_of_vars_found_in_defined_vars"] = found_defined_var_count
    d["num_of_vars_covered_by_context"] = covered_var_count
    return d

def get_file_type(yaml_before_task):
    if yaml_before_task == "":
        file_type = "taskfile"
        context_data = []
        return file_type, context_data

    file_type = ""
    try:
        context_data = yaml.safe_load(yaml_before_task)
    except Exception as e:
        print('the received context could not be loaded as a YAML', e)
        return file_type, []

    if isinstance(context_data, list):
        file_type = "taskfile"
        if context_data[0] and any(
            play_keyword in context_data[0]
                for play_keyword in ["tasks", "pre_tasks", "post_tasks", "handlers", "hosts", "roles"]
            ):
            file_type = "playbook"
    return file_type, context_data


def make_task_format_context(file_type, context_data, defined_vars, parents):
    context_yml = ""
    updated = False

    if not file_type:
        context_yml = yaml.dump(context_data, sort_keys=False)
        return context_yml, updated

    if file_type == "taskfile":
        context_tasks = []
        if defined_vars:
            task = {}
            task["name"] = "define variables"
            task["ansible.builtin.set_fact"] = make_context_vars_dict(defined_vars)
            context_tasks.append(task)
            updated = True

        import_context_tasks = make_import_context(parents)
        if len(import_context_tasks) != 0:
            updated = True
            context_tasks.extend(import_context_tasks)
        context_tasks.extend(context_data)
        # context_yml = yaml.dump(context_tasks, sort_keys=False)
        context_yml = ruamel.yaml.dump(context_tasks, Dumper=ruamel.yaml.RoundTripDumper)

    elif file_type == "playbook":
        context_tasks = make_import_context(parents)
        play = context_data[-1]
        if defined_vars:
            if "vars" not in play:
                play["vars"] = make_context_vars_dict(defined_vars)
                updated = True
            elif "vars" in play:
                original_vars = play["vars"]
                original_vars.update(make_context_vars_dict(defined_vars))
                play["vars"] = original_vars
                updated = True
        if len(context_tasks) > 0:
            for play_keyword in ["tasks", "pre_tasks", "post_tasks", "handlers"]:
                if play_keyword in play:
                    tasks = play[play_keyword]
                    if len(context_tasks) > 0:
                        updated = True
                    if type(tasks) == list:
                        context_tasks.extend(tasks)
                    play[play_keyword] = context_tasks
                    break
        context_data[-1] = play
        # context_yml = yaml.dump(context_data, sort_keys=False)
        context_yml = ruamel.yaml.dump(context_data, Dumper=ruamel.yaml.RoundTripDumper)
    # print("dense context\n", context_yml)
    return context_yml, updated

def make_context_vars_dict(defined_vars):
    vars_dict = {}
    for dv, value in defined_vars.items():
        if value is not None and value != "":
            vars_dict[dv] = value
        else:
            dv = dv.replace('"','')
            dv = dv.replace("'","")
            vars_dict[dv] = f"{{{{ {dv} }}}}"
    return vars_dict

def make_import_context(parents):
    context_tasks = []
    if parents["role"]:
        for pr in parents["role"]:
            task = {}
            if pr.get("description", ""):
                task["name"] = pr["description"]
            else:
                task["name"] = "import role"
            task["ansible.builtin.import_role"] = {
                "name": pr["name"],
            }
            context_tasks.append(task)
    if parents["taskfile"]:
        for t_name in parents["taskfile"]:
            task = {}
            task["name"] = "import tasks"
            task["ansible.builtin.import_tasks"] = {
                "file": t_name,
            }
            context_tasks.append(task)
    return context_tasks

def get_file_path(task, scan_type, scan_path):
    path = ""
    defined_in = getattr(task.spec, "defined_in", "")
    if scan_type == "taskfile":
        path = scan_path
    else:
        path = os.path.join(scan_path, defined_in)
    return path

@dataclass
class PreProcessingRule(Rule):
    rule_id: str = "PP006"
    description: str = "get a task list including all context information"
    enabled: bool = True
    name: str = "GetAnsibleStructureContext"
    version: str = "v0.0.1"
    severity: Severity = Severity.NONE
    tags: tuple = ("wisdom")
    precedence: int = 1

    _char_remove_list: tuple = ("[","]","'",'"','(',')')

    task_context_data_save_filepath: str = default_task_context_data_filepath
    scan_result_save_filepath: str = default_scan_res_filepath

    data_buffer_ftdata: list = field(default_factory=list)
    data_buffer_scan_result: list = field(default_factory=list)
    buffer_size: int = 1000

    def __post_init__(self):

        if os.path.exists(self.task_context_data_save_filepath):
            os.remove(self.task_context_data_save_filepath)

        if os.path.exists(self.scan_result_save_filepath):
            os.remove(self.scan_result_save_filepath)

    def save_data(self):
        # ftdata
        filepath = self.task_context_data_save_filepath

        dirpath = os.path.dirname(filepath)
        if not os.path.exists(dirpath):
            os.makedirs(name=dirpath, mode=0o777, exist_ok=True)

        with open(filepath, "a+") as file:
            for line in self.data_buffer_ftdata:
                file.write(line + "\n")

        self.data_buffer_ftdata = []

        # scan result
        if do_save_scan_result:
            filepath = self.scan_result_save_filepath

            dirpath = os.path.dirname(filepath)
            if not os.path.exists(dirpath):
                os.makedirs(name=dirpath, mode=0o777, exist_ok=True)

            with open(filepath, "a+") as file:
                for line in self.data_buffer_scan_result:
                    file.write(line + "\n")

            self.data_buffer_scan_result = []
        return

    def match(self, ctx: AnsibleRunContext) -> bool:
        return ctx.current.type == RunTargetType.Task

    def process(self, ctx: AnsibleRunContext):
        task = ctx.current

        verdict = False
        detail = {}

        parent_of_task = task.spec.collection or task.spec.role

        if "." in parent_of_task:
            context_parent = ""
            if isinstance(ctx.parent, Collection):
                context_parent = ctx.parent.name
            elif isinstance(ctx.parent, Role):
                context_parent = ctx.parent.fqcn
            elif isinstance(ctx.parent, Repository):
                context_parent = ctx.parent.name
            if context_parent != parent_of_task:
                return RuleResult(verdict=verdict, detail=detail, file=task.file_info(), rule=self.get_metadata())

        verdict = True
        detail = {}

        current_id = task.node_id  # e.g) 0.1.2
        current_depth = task.depth  # e.g) 3

        module = get_module_name(task)
        used_vars, _ = get_used_vars(ctx, task)
        defined_vars, _ = get_defined_vars(task)

        parents = get_parents(ctx, task)
        defined_var_name_values = get_defined_vars_in_parent(ctx, task)

        yaml_before_task = make_yaml_before_task(ctx, task.spec)
        file_type, context_data = get_file_type(yaml_before_task)
        ari_new_context, context_updated = make_task_format_context(file_type, context_data, defined_var_name_values, parents)

        train = {
            "license": "",
            "license_check": "",
            "source": "",
            "path": "",
            "repo_name": "",
            "type": "",
            "prompt": "",
            "input_script": "",
            "metrics": {},
            "output_script": "",
            "token_count": 0,
            "op_token_count": 0,
            "sample_type": 0,
            "context_len": 0,
            "module_name": "",
            "id": 0,
            "prompt_key": ""
        }

        # repo_type, repo_name, license, source, path = find_repo_info(ctx)
        scan_data = ctx.scan_metadata
        scan_type = scan_data.get("type")
        scan_path = scan_data.get("name")
        path = get_file_path(task, scan_type, scan_path)

        # input_script = json.dumps(wisdom_context, default=serialize)
        train["input_script"] = yaml_before_task
        if not context_updated:
            ari_new_context = yaml_before_task
        train["ari_new_context"] = ari_new_context
        train["is_context_updated"] = context_updated
        train["context_len"] = len(ari_new_context)
        # train["yaml_before_task"] = yaml_before_task
        train["license"] = ""
        train["license_check"] = ""
        train["source"] = ""
        train["path"] = getattr(task.spec, "defined_in", "")
        train["repo_name"] = ""
        train["type"] = file_type
        train["scan_type"] = getattr(task.spec, "type", "")

        need_correction = task.get_annotation(key="module.need_correction")
        if need_correction and "." not in task.spec.module:
            content = task.content
            content.set_module_name(module)
            task_str = content.yaml()
        else:
            task_str = getattr(task.spec, "yaml_lines", "")
            if "<<:" in task_str:
                task_str = task.spec.yaml(use_yaml_lines=False)
        train["output_script"] = task_str
        train["module_name"] = module
        train["id"] = f"{current_id}-{current_depth}"  # <ari_node_id> - <ari_node_depth>
        train["ari_task_key"] = task.spec.key
        train["scan_type"] = scan_type
        train["scan_path"] = scan_path

        prompt = ""
        yamllines = getattr(task.spec, "yaml_lines", "").split("\n")
        for yl in yamllines:
            yl2 = yl.replace(" ", "")
            if yl2.startswith("-name:"):
                prompt = yl
                break
        train["prompt"] = prompt.lower()
        lp = prompt.replace("- name:", "").lower()
        train["prompt_key"] = re.sub(r"[^a-zA-Z]", "", lp)

        detail = train

        task_spec = vars(task.spec)
        task_spec["scan_path"] = scan_path
        annotations = {}
        annotations["correct_fqcn"] = task.get_annotation(key="module.correct_fqcn")
        annotations["need_correction"] = task.get_annotation(key="module.need_correction")
        annotations["wrong_module_name"] = task.get_annotation(key="module.wrong_module_name")
        annotations["suggested_fqcn"] = task.get_annotation(key="module.suggested_fqcn")
        annotations["suggested_dependency"] = task.get_annotation(key="module.suggested_dependency")
        annotations["resolved_fqcn"] = task.get_annotation(key="module.resolved_fqcn")
        annotations["not_exist"] = task.get_annotation(key="module.not_exist")
        annotations["wrong_arg_keys"] = task.get_annotation(key="module.wrong_arg_keys")
        annotations["available_arg_keys"] = task.get_annotation(key="module.available_arg_keys")
        annotations["required_arg_keys"] = task.get_annotation(key="module.required_arg_keys")
        annotations["used_alias_and_real_keys"] = task.get_annotation(key="module.used_alias_and_real_keys")
        annotations["wrong_arg_values"] = task.get_annotation(key="module.wrong_arg_values")
        annotations["undefined_values"] = task.get_annotation(key="module.undefined_values")
        annotations["unknown_type_values"] = task.get_annotation(key="module.unknown_type_values")
        annotations["undefined_vars"] = task.get_annotation(key="module.undefined_vars")
        annotations["unknown_name_vars"] = task.get_annotation(key="module.unknown_name_vars")
        annotations["unnecessary_loop_vars"] = task.get_annotation(key="module.unnecessary_loop_vars")
        annotations["vars_used_in_previous_tasks"] = used_vars
        annotations["vars_defined_in_previous_tasks"] = defined_vars
        annotations["vars_used_in_target_task"] = get_used_vars_in_target_task(task)

        task_spec["annotations"] = annotations

        is_last_node_for_this_tree = False
        if ctx.is_last_task(task):
            is_last_node_for_this_tree = True

        self.data_buffer_ftdata.append(json.dumps(detail, default=serialize))
        self.data_buffer_scan_result.append(jsonpickle.encode(task_spec, make_refs=False, unpicklable=False))
        if len(self.data_buffer_ftdata) >= self.buffer_size or is_last_node_for_this_tree:
            self.save_data()

        return RuleResult(verdict=verdict, detail=detail, file=task.file_info(), rule=self.get_metadata())


# NOTE: This is irreversible serialization
def serialize(obj):
    if isinstance(obj, date):
        serial = obj.isoformat()
        return serial

    if isinstance(obj, time):
        serial = obj.isoformat()
        return serial

    if isinstance(obj, bytes):
        return str(obj)

    if hasattr(obj, "__dict__"):
        return obj.__dict__

    return obj
sage_scan/custom_scan/rules/PP006_preprocess_rule_ansible_structure_context.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass, field
from typing import List, Dict
import logging
import json
import jsonpickle
from ansible_risk_insight.models import (
    Module as ARIModule,
    Task as ARITask,
    TaskFile as ARITaskFile,
    Role as ARIRole,
    Playbook as ARIPlaybook,
    Play as ARIPlay,
    Collection as ARICollection,
    Repository as ARIRepository,
    Annotation,
    BecomeInfo,
)
from ansible_risk_insight.findings import Findings as ARIFindings
from ansible_risk_insight.keyutil import get_obj_type, key_delimiter


logger = logging.getLogger(__name__)


@dataclass
class SageObject(object):
    type: str = ""
    key: str = ""
    name: str = ""
    source: dict = field(default_factory=dict)
    source_id: str = ""
    test_object: bool = False
    annotations: Dict[str, any] = field(default_factory=dict)

    @classmethod
    def from_ari_obj(cls, ari_obj, source: dict={}):
        sage_obj = cls()
        if not hasattr(ari_obj, "__dict__"):
            return sage_obj

        attr_mapping = {
            "defined_in": "filepath",
            "path": "filepath",
            "source": "ari_source",
        }
        for key, val in ari_obj.__dict__.items():
            attr_name = key
            if key in attr_mapping:
                attr_name = attr_mapping[key]
            
            if hasattr(sage_obj, attr_name):
                setattr(sage_obj, attr_name, val)

        type_mapping = {
            "repository": "project",
        }
        if sage_obj.type in type_mapping:
            sage_obj.type = type_mapping[sage_obj.type]

        sage_obj.set_source(source)
        return sage_obj
    
    @classmethod
    def to_ari_obj(cls, sage_obj):
        type_str = cls.__name__.lower()

        cls_mapping = {
            "module": ARIModule,
            "task": ARITask,
            "taskfile": ARITaskFile,
            "role": ARIRole,
            "playbook": ARIPlaybook,
            "play": ARIPlay,
            "collection": ARICollection,
            "project": ARIRepository,
        }
        cls = cls_mapping.get(type_str, None)
        if not cls:
            return None
        
        ari_obj = cls()
        if not hasattr(sage_obj, "__dict__"):
            return ari_obj

        default_attr_mapping = {
            "filepath": "defined_in",
        }
        type_specific_attr_mapping = {
            "role": {
                "ari_source": "source",
            },
            "collection": {
                "filepath": "path",
            }
        }
        
        attr_mapping = default_attr_mapping.copy()
        if type_str:
            _attr_mapping = type_specific_attr_mapping.get(type_str, {})
            attr_mapping.update(_attr_mapping)

        for key, val in sage_obj.__dict__.items():
            attr_name = key
            if key in attr_mapping:
                attr_name = attr_mapping[key]
            
            if hasattr(ari_obj, attr_name):
                setattr(ari_obj, attr_name, val)

        type_mapping = {
            "project": "repository",
        }
        if ari_obj.type in type_mapping:
            ari_obj.type = type_mapping[ari_obj.type]

        return ari_obj
    
    def set_source(self, source: dict={}):
        self.source = source
        if source:
            self.source_id = json.dumps(source, separators=(',', ':'))
        return
    
    def set_annotation(self, key: str, value: any):
        self.annotations[key] = value
        return

    def get_annotation(self, key: str, __default: any = None):
        return self.annotations.get(key, __default)
    
    def delete_annotation(self, key: str):
        if key in self.annotations:
            self.annotations.pop(key)
        return


@dataclass
class Module(SageObject):
    type: str = "module"
    key: str = ""
    name: str = ""
    fqcn: str = ""
    collection: str = ""
    role: str = ""
    documentation: str = ""
    examples: str = ""
    arguments: list = field(default_factory=list)
    filepath: str = ""
    builtin: bool = False


@dataclass
class Task(SageObject):
    type: str = "task"
    key: str = ""
    name: str = ""
    
    module: str = ""
    index: int = -1
    play_index: int = -1
    filepath: str = ""
    
    role: str = ""
    collection: str = ""
    become: BecomeInfo = None
    variables: dict = field(default_factory=dict)
    module_defaults: dict = field(default_factory=dict)
    registered_variables: dict = field(default_factory=dict)
    set_facts: dict = field(default_factory=dict)
    loop: dict = field(default_factory=dict)
    options: dict = field(default_factory=dict)
    module_options: dict = field(default_factory=dict)
    executable: str = ""
    executable_type: str = ""
    collections_in_play: list = field(default_factory=list)

    yaml_lines: str = ""
    line_num_in_file: list = field(default_factory=list)  # [begin, end]

    # FQCN for Module and Role. Or a file path for TaskFile.  resolved later
    resolved_name: str = ""
    # candidates of resovled_name
    possible_candidates: list = field(default_factory=list)

    # embed these data when module/role/taskfile are resolved
    module_info: dict = field(default_factory=dict)
    include_info: dict = field(default_factory=dict)


@dataclass
class TaskFile(SageObject):
    type: str = "taskfile"
    key: str = ""
    name: str = ""
    filepath: str = ""
    tasks: list = field(default_factory=list)
    role: str = ""
    collection: str = ""
    yaml_lines: str = ""
    variables: dict = field(default_factory=dict)
    module_defaults: dict = field(default_factory=dict)
    options: dict = field(default_factory=dict)
    task_loading: dict = field(default_factory=dict)


@dataclass
class Role(SageObject):
    type: str = "role"
    key: str = ""
    name: str = ""
    filepath: str = ""
    fqcn: str = ""
    metadata: dict = field(default_factory=dict)
    collection: str = ""
    playbooks: list = field(default_factory=list)
    taskfiles: list = field(default_factory=list)
    handlers: list = field(default_factory=list)
    modules: list = field(default_factory=list)
    dependency: dict = field(default_factory=dict)
    requirements: dict = field(default_factory=dict)
    ari_source: str = ""  # collection/scm repo/galaxy

    default_variables: dict = field(default_factory=dict)
    variables: dict = field(default_factory=dict)
    # key: loop_var (default "item"), value: list/dict of item value
    loop: dict = field(default_factory=dict)
    options: dict = field(default_factory=dict)


@dataclass
class Playbook(SageObject):
    type: str = "playbook"
    key: str = ""
    name: str = ""
    filepath: str = ""
    yaml_lines: str = ""
    role: str = ""
    collection: str = ""
    plays: list = field(default_factory=list)
    variables: dict = field(default_factory=dict)
    options: dict = field(default_factory=dict)


@dataclass
class Play(SageObject):
    type: str = "play"
    key: str = ""
    name: str = ""
    filepath: str = ""
    index: int = -1
    role: str = ""
    collection: str = ""
    import_module: str = ""
    import_playbook: str = ""
    pre_tasks: list = field(default_factory=list)
    tasks: list = field(default_factory=list)
    post_tasks: list = field(default_factory=list)
    # not actual Role, but RoleInPlay defined in this playbook
    roles: list = field(default_factory=list)
    module_defaults: dict = field(default_factory=dict)
    options: dict = field(default_factory=dict)
    collections_in_play: list = field(default_factory=list)
    become: BecomeInfo = None
    variables: dict = field(default_factory=dict)

    task_loading: dict = field(default_factory=dict)


@dataclass
class Collection(SageObject):
    type: str = "collection"
    name: str = ""
    key: str = ""
    filepath: str = ""
    metadata: dict = field(default_factory=dict)
    meta_runtime: dict = field(default_factory=dict)
    files: dict = field(default_factory=dict)
    playbooks: list = field(default_factory=list)
    taskfiles: list = field(default_factory=list)
    roles: list = field(default_factory=list)
    modules: list = field(default_factory=list)
    dependency: dict = field(default_factory=dict)
    requirements: dict = field(default_factory=dict)
    variables: dict = field(default_factory=dict)
    options: dict = field(default_factory=dict)


@dataclass
class Project(SageObject):
    type: str = "project"
    key: str = ""
    name: str = ""
    filepath: str = ""
    # if set, this repository is a collection repository
    my_collection_name: str = ""
    playbooks: list = field(default_factory=list)
    roles: list = field(default_factory=list)
    # for playbook scan
    target_playbook_path: str = ""
    # for taskfile scan
    target_taskfile_path: str = ""
    requirements: dict = field(default_factory=dict)
    installed_collections_path: str = ""
    installed_collections: list = field(default_factory=list)
    installed_roles_path: str = ""
    installed_roles: list = field(default_factory=list)
    modules: list = field(default_factory=list)
    taskfiles: list = field(default_factory=list)
    inventories: list = field(default_factory=list)
    version: str = ""

    @classmethod
    def from_ari_obj(cls, ari_obj, source: dict={}):
        instance = super().from_ari_obj(ari_obj, source)
        instance.key = f"project {instance.source_id}"
        return instance


def convert_to_sage_obj(ari_obj, source: dict={}):
    if isinstance(ari_obj, ARIModule):
        return Module.from_ari_obj(ari_obj, source)
    elif isinstance(ari_obj, ARITask):
        return Task.from_ari_obj(ari_obj, source)
    elif isinstance(ari_obj, ARITaskFile):
        return TaskFile.from_ari_obj(ari_obj, source)
    elif isinstance(ari_obj, ARIRole):
        return Role.from_ari_obj(ari_obj, source)
    elif isinstance(ari_obj, ARIPlaybook):
        return Playbook.from_ari_obj(ari_obj, source)
    elif isinstance(ari_obj, ARIPlay):
        return Play.from_ari_obj(ari_obj, source)
    elif isinstance(ari_obj, ARICollection):
        return Collection.from_ari_obj(ari_obj, source)
    elif isinstance(ari_obj, ARIRepository):
        return Project.from_ari_obj(ari_obj, source)
    else:
        raise ValueError(f"{type(ari_obj)} is not a supported type for Sage objects")


def convert_to_ari_obj(sage_obj):
    if isinstance(sage_obj, Module):
        return Module.to_ari_obj(sage_obj)
    elif isinstance(sage_obj, Task):
        return Task.to_ari_obj(sage_obj)
    elif isinstance(sage_obj, TaskFile):
        return TaskFile.to_ari_obj(sage_obj)
    elif isinstance(sage_obj, Role):
        return Role.to_ari_obj(sage_obj)
    elif isinstance(sage_obj, Playbook):
        return Playbook.to_ari_obj(sage_obj)
    elif isinstance(sage_obj, Play):
        return Play.to_ari_obj(sage_obj)
    elif isinstance(sage_obj, Collection):
        return Collection.to_ari_obj(sage_obj)
    elif isinstance(sage_obj, Project):
        return Project.to_ari_obj(sage_obj)
    else:
        raise ValueError(f"{type(sage_obj)} is not a supported type for ARI objects")


attr_list = [
    "collections",
    "modules",
    "playbooks",
    "plays",
    "projects",
    "roles",
    "taskfiles",
    "tasks",
]


@dataclass
class SageProject(object):
    source: dict = field(default_factory=dict)
    source_id: str = ""

    yml_files: list = field(default_factory=list)

    collections: list = field(default_factory=list)
    modules: list = field(default_factory=list)
    playbooks: list = field(default_factory=list)
    plays: list = field(default_factory=list)
    projects: list = field(default_factory=list)
    roles: list = field(default_factory=list)
    taskfiles: list = field(default_factory=list)
    tasks: list = field(default_factory=list)

    path: str = ""
    scan_timestamp: str = ""
    scan_time_detail: list = field(default_factory=list)
    dir_size: int = 0
    pipeline_version: str = ""

    # only used for ARI RAM conversion
    ari_metadata: dict = field(default_factory=dict)
    dependencies: list = field(default_factory=list)

    @classmethod
    def from_source_objects(
        cls,
        source: dict,
        yml_inventory: list,
        objects: list,
        metadata: dict,
        scan_time: list,
        dir_size: int,
        ari_metadata: dict={},
        dependencies: list=[],
        ):
        proj = cls()
        proj.source = source
        if source:
            proj.source_id = json.dumps(source, separators=(',', ':'))

        proj.yml_files = yml_inventory

        for obj in objects:
            proj.add_object(obj)

        proj.path = metadata.get("name", "")
        proj.scan_timestamp = metadata.get("scan_timestamp", "")
        proj.pipeline_version = metadata.get("pipeline_version", "")
        proj.scan_time_detail = scan_time
        proj.dir_size = dir_size

        proj.ari_metadata = ari_metadata
        proj.dependencies = dependencies
        return proj
    
    def add_object(self, obj: SageObject):
        obj_type = obj.type + "s"
        objects_per_type = getattr(self, obj_type, [])
        objects_per_type.append(obj)
        setattr(self, obj_type, objects_per_type)
        return

    def get_object(self, key: str=""):
        if key:
            obj_type = get_obj_type(key) + "s"
            objects_per_type = getattr(self, obj_type, [])
            for obj in objects_per_type:
                if obj.key == key:
                    return obj
        return None

    def get_all_call_sequences(self, follow_include: bool=True):
        found_taskfile_keys = set()
        all_call_sequences = []
        for p in self.playbooks:
            call_graph = self._get_call_graph(obj=p, follow_include=follow_include)
            call_seq = self.call_graph2sequence(call_graph)
            all_call_sequences.append(call_seq)
            tmp_taskfile_keys = set([obj.key for obj in call_seq if isinstance(obj, TaskFile)])
            found_taskfile_keys = found_taskfile_keys.union(tmp_taskfile_keys)
        for r in self.roles:
            call_graph = self._get_call_graph(obj=r, follow_include=follow_include)
            call_seq = self.call_graph2sequence(call_graph)
            all_call_sequences.append(call_seq)
            tmp_taskfile_keys = set([obj.key for obj in call_seq if isinstance(obj, TaskFile)])
            found_taskfile_keys = found_taskfile_keys.union(tmp_taskfile_keys)
        for tf in self.taskfiles:
            if tf.key in found_taskfile_keys:
                continue
            call_graph = self._get_call_graph(obj=tf, follow_include=follow_include)
            call_seq = self.call_graph2sequence(call_graph)
            all_call_sequences.append(call_seq)
        return all_call_sequences
    
    # NOTE: currently this returns only 1 sequence found first
    def get_call_sequence_for_task(self, task: Task, follow_include: bool=True):
        target_key = task.key
        all_call_seqs = self.get_all_call_sequences(follow_include=follow_include)
        found_seq = None
        for call_seq in all_call_seqs:
            keys_in_seq = [obj.key for obj in call_seq]
            if target_key in keys_in_seq:
                found_seq = call_seq
                break
        return found_seq
    
    def get_call_sequence_by_entrypoint(self, entrypoint: Playbook|Role|TaskFile, follow_include: bool=True):
        call_tree = self.get_call_tree_by_entrypoint(entrypoint=entrypoint, follow_include=follow_include)
        call_seq = self.call_graph2sequence(call_tree)
        return call_seq

    def get_call_tree_by_entrypoint(self, entrypoint: Playbook|Role|TaskFile, follow_include: bool=True):
        return self._get_call_graph(obj=entrypoint, follow_include=follow_include)
    
    def call_graph2sequence(self, call_graph: list=[]):
        if not call_graph:
            return []
        call_seq = [call_graph[0][0]]
        for (_, c_obj) in call_graph:
            call_seq.append(c_obj)
        return call_seq

    # get call graph which starts from the specified object (e.g. playbook -> play -> task)
    def _get_call_graph(self, obj: SageObject=None, key: str="", follow_include: bool=True):
        if not obj and not key:
            raise ValueError("either `obj` or `key` must be non-empty value")
        
        if not obj and key:
            obj = self.get_object(key)
            if not obj:
                raise ValueError(f"No object found for key `{key}`")
        
        history = []
        return self._recursive_get_call_graph(obj=obj, history=history, follow_include=follow_include)


    def _get_children_keys_for_graph(self, obj, follow_include: bool=True):
        if isinstance(obj, Playbook):
            return obj.plays
        elif isinstance(obj, Role):
            # TODO: support role invokation for non main.yml
            target_filenames = ["main.yml", "main.yaml"]
            def get_filename(tf_key):
                return tf_key.split(key_delimiter)[-1].split("/")[-1]
            taskfile_key = [
                tf_key
                for tf_key in obj.taskfiles if get_filename(tf_key) in target_filenames
            ]
            return taskfile_key
        elif isinstance(obj, Play):
            roles = []
            if follow_include:
                if obj.roles:
                    for rip in obj.roles:
                    
                        role_key = rip.role_info.get("key", None)
                        if role_key:
                            roles.append(role_key)
            return obj.pre_tasks + obj.tasks + roles + obj.post_tasks
        elif isinstance(obj, TaskFile):
            return obj.tasks
        elif isinstance(obj, Task):
            if follow_include:
                if obj.include_info:
                    c_key = obj.include_info.get("key", None)
                    if c_key:
                        return [c_key]
        
        return []
        
    def _recursive_get_call_graph(self, obj, history=None, follow_include: bool=True):
        if not history:
            history = []
        
        obj_key = obj.key
        if obj_key in history:
            return []
        
        _history = [h for h in history]
        _history.append(obj_key)

        call_graph = []
        children_keys = self._get_children_keys_for_graph(obj=obj, follow_include=follow_include)
        if children_keys:
            for c_key in children_keys:
                c_obj = self.get_object(c_key)
                if not c_obj:
                    logger.warn(f"No object found for key `{c_key}`; skip this node")
                    continue
                call_graph.append((obj, c_obj))
                sub_graph = self._recursive_get_call_graph(obj=c_obj, history=_history, follow_include=follow_include)
                if sub_graph:
                    call_graph.extend(sub_graph)
        return call_graph
    
    def object_to_key(self):
        
        new_proj = SageProject(
            source=self.source,
            source_id=self.source_id,
            yml_files=self.yml_files,
            path=self.path,
            scan_timestamp=self.scan_timestamp,
        )
        for attr in attr_list:
            objects = getattr(self, attr, [])
            keys = [obj.key for obj in objects]
            setattr(new_proj, attr, keys)
        return new_proj
    
    def metadata(self):
        objects = {}
        for attr in attr_list:
            objects_per_type = getattr(self, attr, [])
            objects[attr] = len(objects_per_type)
        return {
            "source": self.source,
            "source_id": self.source_id,
            "objects": objects,
            "path": self.path,
            "scan_timestamp": self.scan_timestamp,
            "scan_time_detail": self.scan_time_detail,
            "dir_size": self.dir_size,
            "pipeline_version": self.pipeline_version,
            "yml_files": self.yml_files,
            "ari_metadata": self.ari_metadata,
            "dependencies": self.dependencies,
        }
    
    def objects(self):
        _objects = []
        for attr in attr_list:
            objects_per_type = getattr(self, attr, [])
            _objects.extend(objects_per_type)
        return _objects


@dataclass
class SageObjects(object):
    _projects: List[SageProject] = field(default_factory=list)

    def projects(self):
        return self._projects
    
    def project(self, source_type: str="", repo_name: str=""):
        for proj in self.projects():
            p_source_type = proj.source.get("type", "")
            p_repo_name = proj.source.get("repo_name", "")
            if p_source_type == source_type and p_repo_name == repo_name:
                return proj
        return None


def load_objects(fpath: str) -> SageObjects:
    proj_dict = {}
    with open(fpath, "r") as file:
        for line in file:
            obj = jsonpickle.decode(line)
            if not isinstance(obj, SageObject):
                raise ValueError(f"expected type: SageObject, detected type: {type(obj)}")
            source = obj.source
            source_id = obj.source_id
            if source_id not in proj_dict:
                proj_dict[source_id] = SageProject(source=source, source_id=source_id)
            proj_dict[source_id].add_object(obj)
    
    proj_list = [c for c in proj_dict.values()]
    obj = SageObjects(_projects=proj_list)
    return obj


def save_objects(fpath: str, sage_objects: SageObjects):
    projects = sage_objects.projects()
    all_objects = []
    for project in projects:
        _objects = project.objects()
        all_objects.extend(_objects)
    lines = []
    for obj in all_objects:
        if not isinstance(obj, SageObject):
            raise ValueError(f"expected type: SageObject, detected type: {type(obj)}")
        line = jsonpickle.encode(obj, make_refs=False, separators=(',', ':')) + "\n"
        lines.append(line)
    with open(fpath, "w") as file:
        file.write("".join(lines))
    return


# inherit this class to define custom metrics
@dataclass
class DataMetrics(object):
    # sample metrics
    num_of_tasks: int = None
    num_of_plays: int = None


@dataclass
class PlaybookData(object):
    object: Playbook = None
    project: SageProject = None

    # attributes below are automatically set by __post_init__() 

    # parent if any
    collection: Collection = None
    role: Role = None

    # children / include targets
    playbooks: list = field(default_factory=list)
    plays: list = field(default_factory=list)
    roles: list = field(default_factory=list)
    taskfiles: list = field(default_factory=list)
    tasks: list = field(default_factory=list)

    call_tree: list = field(default_factory=list)
    call_seq: list = field(default_factory=list)
    # initialize this with None to clarify whether it is already computed or not
    used_vars: dict = None
    follow_include_for_used_vars: bool = True

    metrics: DataMetrics = field(default_factory=DataMetrics)

    def __post_init__(self):
        if not self.object:
            return
        
        if not self.project:
            return
        
        if self.object.collection:
            for coll in self.project.collections:
                if coll.fqcn == self.object.collection:
                    self.collection = coll
                    break

        if self.object.role:
            for role in self.project.roles:
                if role.fqcn == self.object.role:
                    self.role = role
                    break

        call_tree = self.project.get_call_tree_by_entrypoint(self.object)
        self.call_tree = call_tree

        call_seq = self.project.call_graph2sequence(call_tree)
        self.call_seq = call_seq

        for obj in call_seq:
            if isinstance(obj, Playbook):
                self.playbooks.append(obj)
            elif isinstance(obj, Play):
                self.plays.append(obj)
            elif isinstance(obj, Role):
                self.roles.append(obj)
            elif isinstance(obj, TaskFile):
                self.taskfiles.append(obj)
            elif isinstance(obj, Task):
                self.tasks.append(obj)


    def metrics_to_json(self):
        data = {}
        data["source"] = self.project.source
        data["filepath"] = self.object.filepath
        for k, v in self.metrics.__dict__.items():
            data[k] = v
        return json.dumps(data, separators=(',', ':'))

    def compute_metrics(self, metrics_funcs: list=None):
        # if no func is passed, run sample funcs
        if metrics_funcs is None:
            metrics_funcs = [
                ("num_of_tasks", self.get_num_of_tasks),
                ("num_of_plays", self.get_num_of_plays),
            ]
        # run each funcs
        for _item in metrics_funcs:
            if len(_item) < 2:
                raise ValueError("`metrics_funcs` must be a list of (metrics_key, metrics_func, args)")
            elif len(_item) == 2:
                key = _item[0]
                func = _item[1]
                value = func()
                setattr(self.metrics, key, value)
            elif len(_item) > 2:
                key = _item[0]
                func = _item[1]
                args = _item[2]
                value = func(**args)
                setattr(self.metrics, key, value)
        return
    
    def get_num_of_tasks(self):
        return len(self.get_tasks_in_this_playbook())

    def get_tasks_in_this_playbook(self):
        return [t for t in self.tasks if t.filepath == self.object.filepath]

    def get_num_of_plays(self):
        return len(self.get_plays_in_this_playbook())

    def get_plays_in_this_playbook(self):
        return [p for p in self.plays if p.filepath == self.object.filepath]

@dataclass
class TaskFileData(object):
    object: TaskFile = None
    project: SageProject = None

    # attributes below are automatically set by __post_init__() 

    # parent if any
    collection: Collection = None
    role: Role = None

    # children / include targets
    roles: list = field(default_factory=list)
    taskfiles: list = field(default_factory=list)
    tasks: list = field(default_factory=list)

    call_tree: list = field(default_factory=list)
    call_seq: list = field(default_factory=list)
    # initialize this with None to clarify whether it is already computed or not
    used_vars: dict = None
    follow_include_for_used_vars: bool = True

    metrics: DataMetrics = field(default_factory=DataMetrics)

    def __post_init__(self):
        if not self.object:
            return
        
        if not self.project:
            return
        
        if self.object.collection:
            for coll in self.project.collections:
                if coll.fqcn == self.object.collection:
                    self.collection = coll
                    break

        if self.object.role:
            for role in self.project.roles:
                if role.fqcn == self.object.role:
                    self.role = role
                    break

        call_tree = self.project.get_call_tree_by_entrypoint(self.object)
        self.call_tree = call_tree

        call_seq = self.project.call_graph2sequence(call_tree)
        self.call_seq = call_seq

        for obj in call_seq:
            if isinstance(obj, Role):
                self.roles.append(obj)
            elif isinstance(obj, TaskFile):
                self.taskfiles.append(obj)
            elif isinstance(obj, Task):
                self.tasks.append(obj)


    def metrics_to_json(self):
        data = {}
        data["source"] = self.project.source
        data["filepath"] = self.object.filepath
        for k, v in self.metrics.__dict__.items():
            data[k] = v
        return json.dumps(data, separators=(',', ':'))

    def compute_metrics(self, metrics_funcs: list=None):
        # if no funcs are passed, run sample funcs below
        if metrics_funcs is None:
            metrics_funcs = [
                ("num_of_tasks", self.get_num_of_tasks),
                ("num_of_plays", self.get_num_of_plays),
            ]
        # run each funcs
        for _item in metrics_funcs:
            if len(_item) < 2:
                raise ValueError("`metrics_funcs` must be a list of (metrics_key, metrics_func, args)")
            elif len(_item) == 2:
                key = _item[0]
                func = _item[1]
                value = func()
                setattr(self.metrics, key, value)
            elif len(_item) > 2:
                key = _item[0]
                func = _item[1]
                args = _item[2]
                value = func(**args)
                setattr(self.metrics, key, value)
        return
    
    def get_num_of_tasks(self):
        return len(self.get_tasks_in_this_taskfile())

    def get_tasks_in_this_taskfile(self):
        return [t for t in self.tasks if t.filepath == self.object.filepath]
    
    def get_num_of_plays(self):
        return 0

sage_scan/models.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass, field
import datetime
from ansible_risk_insight.scanner import ARIScanner, config, Config
from ansible_risk_insight.models import NodeResult, RuleResult, AnsibleRunContext, Object
from ansible_risk_insight.finder import (
    find_all_ymls,
    label_yml_file,
    get_role_info_from_path,
    get_project_info_for_file,
)
from ansible_risk_insight.utils import escape_local_path
from sage_scan.utils import get_rule_id_list, get_git_version
from sage_scan.models import convert_to_sage_obj, SageProject
import os
import time
import traceback
import logging
import threading
import json
import jsonpickle


logging.basicConfig()
logger = logging.getLogger("pipeline")
log_level_str = os.getenv("SAGE_LOG_LEVEL", "info")
log_level_map = {
    "error": logging.ERROR,
    "warning": logging.WARNING,
    "info": logging.INFO,
    "debug": logging.DEBUG,
}
log_level = log_level_map.get(log_level_str, None)
if log_level is None:
    logger.warn(f"logging level \"{log_level_str}\" is not supported. Set to \"info\" instead.")
    log_level = logging.INFO
logger.setLevel(log_level)

ftdata_rule_dir = os.path.join(os.path.dirname(__file__), "custom_scan/rules")


@dataclass
class InputData:
    index: int = 0
    total_num: int = 0
    type: str = ""
    name: str = ""
    path: str = ""
    yaml: str = ""
    metadata: dict = field(default_factory=dict)

@dataclass
class OutputData:
    input: InputData = field(default_factory=InputData)
    data: any = None
    metadata: dict = field(default_factory=dict)

    # convert the output data into a json line string
    def serialize(self):
        return jsonpickle.encode(self.data, make_refs=False, separators=(',', ':'))

@dataclass
class SerializableRunContext(object):
    targets: list = field(default_factory=list)
    root_key: str = ""
    parent: Object = None
    scan_metadata: dict = field(default_factory=dict)
    last_item: bool = False

    @classmethod
    def from_ansible_run_context(cls, rc: AnsibleRunContext):
        src = cls()
        src.targets = rc.sequence.items
        src.root_key = rc.root_key
        src.parent = rc.parent
        src.scan_metadata = rc.scan_metadata
        src.last_item = rc.last_item
        return src
    
    def to_ansible_run_context(self):
        rc = AnsibleRunContext.from_targets(
            targets=self.targets,
            root_key=self.root_key,
            parent=self.parent,
            scan_metadata=self.scan_metadata,
            last_item=self.last_item,
        )
        return rc


@dataclass
class SagePipeline(object):
    ari_kb_data_dir: str = ""
    ari_rules_dir: str = ""
    ari_rules: list = field(default_factory=list)

    scanner: ARIScanner = None
    log_level_str: str = ""
    logger: logging.Logger = None

    timeout: float = 0.0

    # whether it scans the failed files later: default to True
    do_multi_stage: bool = True

    # whether it scans the targets in parallel: default to False
    do_parallel: bool = False

    do_save_yml_inventory: bool = True
    do_save_findings: bool = False
    do_save_metadata: bool = True
    do_save_objects: bool = True
    do_save_output: bool = False

    use_ftdata_rule: bool = False 

    aggregation_rule_id: str = ""

    ari_out_dir: str = ""
    ari_include_tests: bool = True
    ari_objects: bool = False

    accumulate: bool = True
    scan_records: dict = field(default_factory=dict)

    # special scan records
    yml_inventory: list = field(default_factory=list)

    def __post_init__(self):
        if not self.logger:
            self.init_logger()

        if not self.scanner:
            self.init_scanner()

    def init_logger(self):
        global logger
        _level_str = self.log_level_str or log_level_str
        _level = log_level_map.get(_level_str, logging.INFO)
        logger.setLevel(_level)
        self.logger = logger

    def init_scanner(self):
        _data_dir = self.ari_kb_data_dir
        if not _data_dir:
            self.logger.debug(f"ARI KB data dir is not configured.")
        _rules_dir = self.ari_rules_dir
        _rule_id_list = []
        if self.use_ftdata_rule:
            if not _rules_dir:
                _rules_dir = ftdata_rule_dir
            _rule_id_list = get_rule_id_list(_rules_dir)
        _rules = []
        if self.ari_rules:
            _rules = self.ari_rules
        else:
            _rules = ["P001", "P002", "P003", "P004"] + _rule_id_list
            self.ari_rules = _rules
            self.aggregation_rule_id = _rules[0] if _rules else ""
        scanner = ARIScanner(
            config=Config(
                data_dir=_data_dir,
                rules_dir=_rules_dir,
                rules=_rules,
                log_level=self.log_level_str,
            ),
            read_ram=False,
            write_ram=False,
            silent=True,
        )
        self.scanner = scanner


    # create a list of InputData to scan
    def to_input(self, **kwargs):
        if isinstance(kwargs, dict) and "target_dir" in kwargs:
            target_dir = kwargs["target_dir"]
            self.scan_records["source"] = kwargs.get("source", {})
            return self._target_dir_to_input(target_dir)
        elif isinstance(kwargs, dict) and "raw_yaml" in kwargs:
            raw_yaml = kwargs["raw_yaml"]
            self.scan_records["single_scan"] = True
            self.scan_records["source"] = kwargs.get("source", {})
            return self._single_yaml_to_input(raw_yaml)
        elif isinstance(kwargs, dict) and "ftdata_path" in kwargs:
            ftdata_path = kwargs["ftdata_path"]
            return self._ftdata_to_input(ftdata_path)
        return
    
    # create a list of OutputData to save as a multi-line json file
    def to_output(self, **kwargs):
        if self.accumulate and "output_dir" in kwargs:
            search_dir = kwargs["output_dir"]
            return self._task_context_to_output(search_dir)
        elif isinstance(kwargs, dict) and "from_aggregated_results" in kwargs:
            return self._aggregated_results_to_output()
        # return SageProject by default
        return self._sage_project_to_output()

    def _target_dir_to_input(self, target_dir):
        dir_size = get_dir_size(target_dir)
        path_list = get_yml_list(target_dir)
        project_file_list, role_file_list, independent_file_list = create_scan_list(path_list)
        # used for detecting missing files at the 1st scan
        self.scan_records["project_file_list"] = project_file_list
        self.scan_records["role_file_list"] = role_file_list
        self.scan_records["independent_file_list"] = independent_file_list
        self.scan_records["non_task_scanned_files"] = []
        self.scan_records["findings"] = []
        self.scan_records["metadata"] = {}
        self.scan_records["time"] = []
        self.scan_records["size"] = dir_size
        self.scan_records["objects"] = []

        num = len(project_file_list) + len(role_file_list) + len(independent_file_list)
        target_counts = []
        if project_file_list:
            target_counts.append(f"{len(project_file_list)} projects")
        if role_file_list:
            target_counts.append(f"{len(role_file_list)} roles")
        if independent_file_list:
            target_counts.append(f"{len(independent_file_list)} playbooks/taskfiles")
        target_str = ", ".join(target_counts)
        total_str = f"(total {len(path_list)} files)"

        self.logger.info(f"Start scanning for {target_str} {total_str}")

        input_list = []

        i = 0
        for project_name in project_file_list:
            project_path = project_file_list[project_name].get("path")
            input_list.append(InputData(
                index=i,
                total_num=num,
                type="project",
                name=project_name,
                path=project_path,
                metadata={
                    "base_dir": project_path,
                }
            ))
            i += 1

        for role_name in role_file_list:
            _type = "role"
            _name = role_name
            role_path = role_file_list[role_name].get("path")
            input_list.append(InputData(
                index=i,
                total_num=num,
                type="role",
                name=role_name,
                path=role_path,
                metadata={
                    "base_dir": role_path,
                }
            ))
            i += 1

        for file in independent_file_list:
            _name = file.get("filepath")
            filepath = _name
            _type = file.get("label")
            if _type in ["playbook", "taskfile"]:
                input_list.append(InputData(
                index=i,
                total_num=num,
                type=_type,
                name=_name,
                path=filepath,
                metadata={
                    "base_dir": target_dir,
                }
            ))
            i += 1

        return input_list

    def _single_yaml_to_input(self, raw_yaml):
        label, _, error = label_yml_file(yml_body=raw_yaml)
        if error:
            raise ValueError(f"failed to detect the input YAML type: {error}")
        if label not in ["playbook", "taskfile"]:
            raise ValueError(f"playbook and taskfile are the only supported types, but the input file is `{label}`")
        input_data = InputData(
            index=0,
            total_num=1,
            yaml=raw_yaml,
            type=label,
        )
        input_list = [input_data]
        return input_list

    # TODO: implement this
    def _ftdata_to_input(self, ftdata_path):
        pass

    def _task_context_to_output(self, search_dir):
        if search_dir is None:
            return []

        prefix = "task_context_data"
        filenames = os.listdir(search_dir)
        output_list = []
        loaded_task_keys = []
        used_task_context_files = []
        for filename in filenames:
            if filename.startswith(prefix):
                fpath = os.path.join(search_dir, filename)
                with open(fpath, "r") as file:
                    for line in file:
                        d = json.loads(line)
                        task_key = d.get("ari_task_key", None)
                        scan_path = d.get("scan_path", None)
                        uniq_task_key = f"{scan_path}-{task_key}"
                        if uniq_task_key in loaded_task_keys:
                            continue
                        if task_key in loaded_task_keys:
                            continue
                        output_data = OutputData(data=d)
                        output_list.append(output_data)
                        loaded_task_keys.append(uniq_task_key)
                used_task_context_files.append(fpath)
        for fpath in used_task_context_files:
            if not fpath:
                continue
            try:
                os.remove(fpath)
            except Exception:
                err = traceback.format_exc()
                self.logger.warn(f"failed to remove the temporary file \"{fpath}\": {err}")
        return output_list
    
    def _sage_project_to_output(self):
        proj = self._create_sage_project()
        output_data = OutputData(data=proj, metadata={"single_object": True})
        output_list = [output_data]
        return output_list

    # TODO: implement this
    def _aggregated_results_to_output(self):
        pass

    def check_timeout(self):
        limit_seconds = self.timeout
        if limit_seconds <= 0:
            return
        
        now = time.time()
        begin = self.scan_records["begin"]
        if (now - begin) > limit_seconds:
            raise ValueError(f"TimeoutError: this scan took more than {limit_seconds} seconds")

    
    def run(self, **kwargs):
        self.logger.info("Running data pipeline")
        
        self._init_scan_records()

        if isinstance(kwargs, dict) and "timeout" in kwargs:
            self.timeout = kwargs["timeout"]
        
        # By default, three types of input are supported
        #   - target_dir: path to a target directory
        #   - raw_yaml: YAML string of a playbook or a taskfile
        #   - (TODO) ftdata_path: path to a ftdata file
        # You can override `to_input` function to define a customized one
        input_list = self.to_input(**kwargs)
        self.check_timeout()

        # Specify output filepath with the following argument
        #   - output_dir: path to the output directory
        output_dir = ""
        if isinstance(kwargs, dict) and "output_dir" in kwargs:
            output_dir = kwargs["output_dir"]

        if isinstance(kwargs, dict) and "scan_func" in kwargs:
            scan_func = kwargs["scan_func"]
            scan_func(input_list)

        # create yml inventory here, but this will be updated after scanning
        self.yml_inventory = self.create_yml_inventory()

        multi_stage = self.do_multi_stage
        if "single_scan" in self.scan_records and self.scan_records["single_scan"]:
            multi_stage = False
        yml_inventory_only = False
        if isinstance(kwargs, dict) and "yml_inventory_only" in kwargs:
            yml_inventory_only = kwargs["yml_inventory_only"]
        if yml_inventory_only:
            self.check_timeout()
            if output_dir and self.do_save_yml_inventory:
                yml_inventory_path = os.path.join(output_dir, "yml_inventory.json")
                self.save_yml_inventory(yml_inventory_path)
                self.check_timeout()
            return

        elif multi_stage:
            self._multi_stage_scan(input_list)
            self.check_timeout()
        else:
            self._single_scan(input_list)
            self.check_timeout()
        
        self.yml_inventory = self.create_yml_inventory()
        if output_dir and self.do_save_yml_inventory:
            yml_inventory_path = os.path.join(output_dir, "yml_inventory.json")
            self.save_yml_inventory(yml_inventory_path)
            self.check_timeout()

        if output_dir and self.do_save_findings:
            findings_path = os.path.join(output_dir, "findings.json")
            self.save_findings(findings_path)
            self.check_timeout()

        if output_dir and self.do_save_metadata:
            metadata_path = os.path.join(output_dir, "sage-metadata.json")
            self.save_metadata(metadata_path)
            self.check_timeout()

        if output_dir and self.do_save_objects:
            objects_path = os.path.join(output_dir, "sage-objects.json")
            self.save_objects(objects_path)
            self.check_timeout()

        output_list = self.to_output(**kwargs)
        self.check_timeout()
        output_data = None
        if output_list:
            if len(output_list) == 1 and output_list[0].metadata.get("single_object"):
                output_data = output_list[0].data
            else:
                output_data = [od.data for od in output_list if isinstance(od, OutputData)]

        self.logger.info("Done")

        self._clear_scan_records()
        
        return output_data

    # TODO: implement this
    def _single_scan(self, input_list):
        start = time.time()
        # first stage scan; scan the input as project
        for input_data in input_list:
            self.scan(start, input_data)
            self.check_timeout()

        self.check_timeout()
        return

    def _multi_stage_scan(self, input_list):

        start = time.time()
        # first stage scan; scan the input as project
        for input_data in input_list:
            self.scan(start, input_data)
            self.check_timeout()

        # make a list of missing files from the first scan
        missing_files = []
        for project_name in self.scan_records["project_file_list"]:
            base_dir = os.path.abspath(self.scan_records["project_file_list"][project_name]["path"])
            for file in self.scan_records["project_file_list"][project_name]["files"]:
                label = file.get("label", "")
                filepath = file.get("filepath", "")
                task_scanned = file.get("task_scanned", False)
                role_info = file.get("role_info", {})
                non_task_scanned = True if filepath in self.scan_records["non_task_scanned_files"] else False
                if not task_scanned and not non_task_scanned and label in ["playbook", "taskfile"]:
                    if role_info and role_info.get("is_external_dependency", False):
                        continue
                    _type = label
                    _name = filepath
                    missing_files.append((_type, _name, filepath, base_dir, "project"))
            self.check_timeout()

        for role_name in self.scan_records["role_file_list"]:
            base_dir = os.path.abspath(self.scan_records["role_file_list"][role_name]["path"])
            for file in self.scan_records["role_file_list"][role_name]["files"]:
                label = file.get("label", "")
                filepath = file.get("filepath", "")
                task_scanned = file.get("task_scanned", False)
                role_info = file.get("role_info", {})
                non_task_scanned = True if filepath in self.scan_records["non_task_scanned_files"] else False
                if not task_scanned and not non_task_scanned and label in ["playbook", "taskfile"]:
                    if role_info and role_info.get("is_external_dependency", False):
                        continue
                    _type = label
                    _name = filepath
                    missing_files.append((_type, _name, filepath, base_dir, "role"))
            self.check_timeout()
        
        self.scan_records["missing_files"] = missing_files
        num_of_missing = len(missing_files)
        second_input_list = [
            InputData(
                index=i,
                total_num=num_of_missing,
                type=_type,
                name=_name,
                path=filepath,
                metadata={"original_type": original_type, "base_dir": base_dir}
            )
            for i, (_type, _name, filepath, base_dir, original_type) in enumerate(missing_files)
        ]
        start = time.time()
        for input_data in second_input_list:
            self.scan(start, input_data)
            self.check_timeout()
        return
    
    def _init_scan_records(self):
        self.scan_records = {
            "project_file_list": {},
            "role_file_list": {},
            "independent_file_list": [],
            "non_task_scanned_files": [],
            "findings": [],
            "metadata": {},
            "time": [],
            "size": 0,
            "objects": [],
            "begin": time.time(),
        }
        self.yml_inventory = []
        return
    
    def _clear_scan_records(self):
        self.scan_records = {}
        return
    
    def scan(self, start, input_data):
        if not isinstance(input_data, InputData):
            raise ValueError(f"input data must be InputData type, but {type(input_data)}")         
        
        i = input_data.index
        num = input_data.total_num
        _type = input_data.type
        name = input_data.name
        path = input_data.path
        raw_yaml = input_data.yaml
        original_type = input_data.metadata.get("original_type", _type)
        base_dir = input_data.metadata.get("base_dir", None)

        kwargs = {
            "type": _type,
        }
        if path:
            kwargs["name"] = path
        if raw_yaml:
            kwargs["raw_yaml"] = raw_yaml
        
        source = self.scan_records.get("source", {})
        display_name = name
        if base_dir and name.startswith(base_dir):
            display_name = name.replace(base_dir, "", 1)
            if display_name and display_name[-1] == "/":
                display_name = display_name[:-1]

        yaml_label_list = []
        if self.yml_inventory:
            for file_info in self.yml_inventory:
                if not isinstance(file_info, dict):
                    continue
                fpath = file_info.get("path_from_root", "")
                label = file_info.get("label", "")
                role_info = file_info.get("role_info", {})
                if not fpath or not label:
                    continue
                yaml_label_list.append((fpath, label, role_info))

        start_of_this_scan = time.time()
        thread_id = threading.get_native_id()
        self.logger.debug(f"[{i+1}/{num}] start {_type} {display_name}")
        use_src_cache = True

        taskfile_only = False
        playbook_only = False
        out_dir_basename = name
        if _type != "role" and _type != "project":
            taskfile_only = True
            playbook_only = True
            out_dir_basename = escape_local_path(name)

        result = None
        scandata = None
        elapsed = None
        try:
            include_tests = self.ari_include_tests
            out_dir = ""
            if self.ari_out_dir:
                out_dir = os.path.join(self.ari_out_dir, _type, out_dir_basename)
            objects = False
            if self.ari_objects and out_dir:
                objects = True
            begin = time.time()
            result = self.scanner.evaluate(
                **kwargs,
                install_dependencies=True,
                include_test_contents=include_tests,
                objects=objects,
                out_dir=out_dir,
                load_all_taskfiles=True,
                use_src_cache=use_src_cache,
                taskfile_only=taskfile_only,
                playbook_only=playbook_only,
                base_dir=base_dir,
                yaml_label_list=yaml_label_list,
            )
            elapsed = time.time() - begin
            scandata = self.scanner.get_last_scandata()
        except Exception:
            error = traceback.format_exc()
            if error:
                self.logger.error(f"Failed to scan {path} in {name}: error detail: {error}")

        if result:
            for target_result in result.targets:
                for node_result in target_result.nodes:
                    if not node_result:
                        continue
                    if not isinstance(node_result, NodeResult):
                        raise ValueError(f"node_result must be a NodeResult instance, but {type(node_result)}")
                    rule_result = node_result.find_result(self.aggregation_rule_id)
                    if not isinstance(rule_result, RuleResult):
                        raise ValueError(f"rule_result must be a RuleResult instance, but {type(rule_result)}")
                    if rule_result.error:
                        spec = node_result.node.spec
                        key = spec.key

        if scandata:
            all_scanned_files = self.get_all_files_from_scandata(scandata, path)
            task_scanned_files = [fpath for fpath, scan_type in all_scanned_files if scan_type == "task"]
            play_scanned_files = [fpath for fpath, scan_type in all_scanned_files if scan_type == "play"]
            found_files = []
            if original_type == "project":
                if name in self.scan_records["project_file_list"]:
                    files_num = len(self.scan_records["project_file_list"][name]["files"])
                    for j in range(files_num):
                        fpath = self.scan_records["project_file_list"][name]["files"][j]["filepath"]
                        if fpath in task_scanned_files:
                            self.scan_records["project_file_list"][name]["files"][j]["task_scanned"] = True
                            self.scan_records["project_file_list"][name]["files"][j]["scanned_as"] = _type
                        elif fpath in play_scanned_files:
                            self.scan_records["project_file_list"][name]["files"][j]["scanned_as"] = _type
                            self.scan_records["non_task_scanned_files"].append(fpath)

            elif original_type == "role":
                if name in self.scan_records["role_file_list"]:
                    files_num = len(self.scan_records["role_file_list"][name]["files"])
                    for j in range(files_num):
                        fpath = self.scan_records["role_file_list"][name]["files"][j]["filepath"]
                        if fpath in task_scanned_files:
                            self.scan_records["role_file_list"][name]["files"][j]["task_scanned"] = True
                            self.scan_records["role_file_list"][name]["files"][j]["scanned_as"] = _type
                        elif fpath in play_scanned_files:
                            self.scan_records["role_file_list"][name]["files"][j]["scanned_as"] = _type
                            self.scan_records["non_task_scanned_files"].append(fpath)
            else:
                files_num = len(self.scan_records["independent_file_list"])
                for j in range(files_num):
                    fpath = self.scan_records["independent_file_list"][j]["filepath"]
                    if fpath in task_scanned_files:
                        self.scan_records["independent_file_list"][j]["task_scanned"] = True
                        self.scan_records["independent_file_list"][j]["scanned_as"] = _type
                    elif fpath in play_scanned_files:
                        self.scan_records["independent_file_list"][j]["scanned_as"] = _type
                        self.scan_records["non_task_scanned_files"].append(fpath)

            findings = scandata.findings
            self.scan_records["findings"].append({"target_type": _type, "target_name": name, "findings": findings})

            trees = scandata.trees
            annotation_dict = {}
            skip_annotation_keys = [
                "",
                "module.available_args",
                "variable.unnecessary_loop_vars",
            ]
            for _tree in trees:
                for call_obj in _tree.items:
                    if not hasattr(call_obj, "annotations"):
                        continue
                    orig_annotations = call_obj.annotations
                    annotations = {anno.key: anno.value for anno in orig_annotations if isinstance(anno.key, str) and anno.key not in skip_annotation_keys}
                    spec_key = call_obj.spec.key
                    if annotations:
                        annotation_dict[spec_key] = annotations

            ari_objects = findings.root_definitions.get("definitions", {})
            for obj_type in ari_objects:
                ari_objects_per_type = ari_objects[obj_type]
                for ari_obj in ari_objects_per_type:
                    ari_spec_key = ari_obj.key
                    sage_obj = convert_to_sage_obj(ari_obj, source)
                    if source:
                        sage_obj.set_source(source)
                    if ari_spec_key in annotation_dict:
                        sage_obj.annotations = annotation_dict[ari_spec_key]
                    self.scan_records["objects"].append(sage_obj)

            self.scan_records["time"].append({"target_type": _type, "target_name": name, "scan_seconds": elapsed})

            if _type == "project":
                metadata = findings.metadata.copy()
                metadata.pop("time_records")
                metadata["scan_timestamp"] = datetime.datetime.utcnow().isoformat(timespec="seconds")
                metadata["pipeline_version"] = get_git_version()
                self.scan_records["metadata"] = metadata

                ari_metadata = findings.metadata.copy()
                dependencies = findings.dependencies.copy()
                self.scan_records["ari_metadata"] = ari_metadata
                self.scan_records["dependencies"] = dependencies

        elapsed_for_this_scan = round(time.time() - start_of_this_scan, 2)
        if elapsed_for_this_scan > 60:
            self.logger.warn(f"It took {elapsed_for_this_scan} sec. to process [{i+1}/{num}] {_type} {name}")


    def get_all_files_from_scandata(self, scandata, scan_root_dir):
        
        task_specs = scandata.root_definitions.get("definitions", {}).get("tasks", [])
        all_files = []
        for task_spec in task_specs:
            fullpath = os.path.join(scan_root_dir, task_spec.defined_in)
            if fullpath not in all_files:
                all_files.append((fullpath, "task"))

        # some plays have only `roles` instead of `tasks`
        # count this type of playbook files here
        play_specs = scandata.root_definitions.get("definitions", {}).get("plays", [])
        for play_spec in play_specs:
            fullpath = os.path.join(scan_root_dir, play_spec.defined_in)
            if fullpath not in all_files:
                all_files.append((fullpath, "play"))
        return all_files
    
    def create_yml_inventory(self):
        yml_inventory = []
        for project_name in self.scan_records["project_file_list"]:
            for file in self.scan_records["project_file_list"][project_name]["files"]:
                task_scanned = file.get("task_scanned", False)
                file["task_scanned"] = task_scanned
                scanned_as = file.get("scanned_as", "")
                file["scanned_as"] = scanned_as
                yml_inventory.append(file)

        for role_name in self.scan_records["role_file_list"]:
            for file in self.scan_records["role_file_list"][role_name]["files"]:
                task_scanned = file.get("task_scanned", False)
                file["task_scanned"] = task_scanned
                scanned_as = file.get("scanned_as", "")
                file["scanned_as"] = scanned_as
                yml_inventory.append(file)

        for file in self.scan_records["independent_file_list"]:
            task_scanned = file.get("task_scanned", False)
            file["task_scanned"] = task_scanned
            scanned_as = file.get("scanned_as", "")
            file["scanned_as"] = scanned_as
            yml_inventory.append(file)
        
        return yml_inventory

    def save_yml_inventory(self, output_path):
        lines = [json.dumps(file) + "\n" for file in self.yml_inventory]

        out_dir = os.path.dirname(output_path)
        if not os.path.exists(out_dir):
            os.makedirs(out_dir, exist_ok=True)
        
        with open(output_path, "w") as outfile:
            outfile.write("".join(lines))

    def save_findings(self, output_path):
        if not self.scan_records:
            return
        if "findings" not in self.scan_records:
            return
        
        findings_list = self.scan_records["findings"]
        lines = []
        for d in findings_list:
            findings = d["findings"]
            findings_json_str = findings.dump()
            lines.append(findings_json_str + "\n")
        
        out_dir = os.path.dirname(output_path)
        if not os.path.exists(out_dir):
            os.makedirs(out_dir, exist_ok=True)

        with open(output_path, "w") as outfile:
            outfile.write("".join(lines))

    def _create_sage_project(self):
        if not self.scan_records:
            return
        source = self.scan_records.get("source", {})
        yml_inventory = self.yml_inventory
        objects = self.scan_records.get("objects", [])
        metadata = self.scan_records.get("metadata", {})
        scan_time = self.scan_records.get("time", [])
        dir_size = self.scan_records.get("size", 0)
        ari_metadata = self.scan_records.get("ari_metadata", {})
        dependencies = self.scan_records.get("dependencies", [])
        proj = SageProject.from_source_objects(
            source=source,
            yml_inventory=yml_inventory,
            objects=objects,
            metadata=metadata,
            scan_time=scan_time,
            dir_size=dir_size,
            ari_metadata=ari_metadata,
            dependencies=dependencies,
        )
        return proj

    def save_metadata(self, output_path):
        if not self.scan_records:
            return
        if "metadata" not in self.scan_records:
            return
        
        proj = self._create_sage_project()        
        proj_metadata = proj.metadata()
        
        out_dir = os.path.dirname(output_path)
        if not os.path.exists(out_dir):
            os.makedirs(out_dir, exist_ok=True)

        with open(output_path, "w") as outfile:
            outfile.write(jsonpickle.encode(proj_metadata, make_refs=False, separators=(',', ':')))

    def save_objects(self, output_path):
        if not self.scan_records:
            return
        if "objects" not in self.scan_records:
            return
        
        objects = self.scan_records["objects"]
        lines = []
        for obj in objects:
            obj_json = jsonpickle.encode(obj, make_refs=False, separators=(',', ':'))
            lines.append(obj_json + "\n")
        
        out_dir = os.path.dirname(output_path)
        if not os.path.exists(out_dir):
            os.makedirs(out_dir, exist_ok=True)

        with open(output_path, "w") as outfile:
            outfile.write("".join(lines))


def get_yml_label(file_path, root_path):
    if root_path and root_path[-1] == "/":
        root_path = root_path[:-1]
    
    relative_path = file_path.replace(root_path, "")
    if relative_path[-1] == "/":
        relative_path = relative_path[:-1]
    
    label, name_count, error = label_yml_file(yml_path=file_path)
    role_name, role_path = get_role_info_from_path(file_path)
    role_info = None
    if role_name and role_path:
        relative_role_path = role_path.replace(root_path, "")
        if relative_role_path[0] == "/":
            relative_role_path = relative_role_path[1:]
        role_info = {"name": role_name, "path": role_path, "relative_path": relative_role_path}

    project_name, project_path = get_project_info_for_file(file_path, root_path)
    project_info = None
    if project_name and project_path:
        project_info = {"name": project_name, "path": project_path}
    
    # print(f"[{label}] {relative_path} {role_info}")
    if error:
        logger.debug(f"failed to get yml label:\n {error}")
        label = "error"
    return label, role_info, project_info, name_count, error


def get_yml_list(root_dir: str):
    found_ymls = find_all_ymls(root_dir)
    all_files = []
    for yml_path in found_ymls:
        label, role_info, project_info, name_count, error = get_yml_label(yml_path, root_dir)
        if not role_info:
            role_info = {}
        if not project_info:
            project_info = {}
        if role_info:
            if role_info["path"] and not role_info["path"].startswith(root_dir):
                role_info["path"] = os.path.join(root_dir, role_info["path"])
            role_info["is_external_dependency"] = True if "." in role_info["name"] else False
        in_role = True if role_info else False
        in_project = True if project_info else False
        all_files.append({
            "filepath": yml_path,
            "path_from_root": yml_path.replace(root_dir, "").lstrip("/"),
            "label": label,
            "role_info": role_info,
            "project_info": project_info,
            "in_role": in_role,
            "in_project": in_project,
            "name_count": name_count,
            "error": error,
        })
    return all_files


def create_scan_list(yml_inventory):
    role_file_list = {}
    project_file_list = {}
    independent_file_list = []
    for yml_data in yml_inventory:
        filepath = yml_data["filepath"]
        path_from_root = yml_data["path_from_root"]
        label = yml_data["label"]
        role_info = yml_data["role_info"]
        in_role = yml_data["in_role"]
        project_info = yml_data["project_info"]
        in_project = yml_data["in_project"]
        name_count = yml_data["name_count"]
        error = yml_data["error"]
        if project_info:
            p_name = project_info.get("name", "")
            p_path = project_info.get("path", "")
            if p_name not in project_file_list:
                project_file_list[p_name] = {"path": p_path, "files": []}
            project_file_list[p_name]["files"].append({
                "filepath": filepath,
                "path_from_root": path_from_root,
                "label": label,
                "project_info": project_info,
                "role_info": role_info,
                "in_project": in_project,
                "in_role": in_role,
                "name_count": name_count,
                "error": error,
            })
        elif role_info:
            r_name = role_info.get("name", "")
            r_path = role_info.get("path", "")
            if role_info.get("is_external_dependency", False):
                continue
            if r_name not in role_file_list:
                role_file_list[r_name] = {"path": r_path, "files": []}
            role_file_list[r_name]["files"].append({
                "filepath": filepath,
                "path_from_root": path_from_root,
                "label": label,
                "project_info": project_info,
                "role_info": role_info,
                "in_project": in_project,
                "in_role": in_role,
                "name_count": name_count,
                "error": error,
            })
        else:
            independent_file_list.append({
                "filepath": filepath,
                "path_from_root": path_from_root,
                "label": label,
                "project_info": project_info,
                "role_info": role_info,
                "in_project": in_project,
                "in_role": in_role,
                "name_count": name_count,
                "error": error,
            })
    return project_file_list, role_file_list, independent_file_list


def get_dir_size(path=""):
    total = 0
    with os.scandir(path) as it:
        for entry in it:
            if entry.is_file(follow_symlinks=False):
                total += entry.stat().st_size
            elif entry.is_dir(follow_symlinks=False):
                total += get_dir_size(entry.path)
    return total
sage_scan/pipeline.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

sage_scan/process/__init__.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from sage_scan.models import (
    Task,
    Module,
)
from ansible_risk_insight.models import (
    ExecutableType,
    Module as ARIModule,
    ActionGroupMetadata,
    VariableType,
    ArgumentsType,
)

ARGUMENTS_ANNOTATION_KEY = "arguments"
VARIABLES_SET_ANNOTATION_KEY = "variables_set"
VARIABLES_USED_ANNOTATION_KEY = "variables_used"
MODULE_OBJECT_ANNOTATION_KEY = "module_object"


# P001 rule in ARI
def set_module_spec_annotations(task: Task):
    resolved_fqcn = ""
    wrong_module_name = ""
    not_exist = False
    correct_fqcn = ""
    need_correction = False
    module = task.get_annotation(MODULE_OBJECT_ANNOTATION_KEY)
    if module:
        resolved_fqcn = module.fqcn
        correct_fqcn = module.fqcn
    else:
        wrong_module_name = task.module
        not_exist = True

    if correct_fqcn != task.module or not_exist:
        need_correction = True

    module_examples = ""
    if module:
        module_examples = module.examples

    task.set_annotation("module.resolved_fqcn", resolved_fqcn)
    task.set_annotation("module.wrong_module_name", wrong_module_name)
    task.set_annotation("module.not_exist", not_exist)
    task.set_annotation("module.correct_fqcn", correct_fqcn)
    task.set_annotation("module.need_correction", need_correction)
    task.set_annotation("module.examples", module_examples)
    return task


# P002 rule in ARI
def set_module_arg_key_annotations(task: Task, knowledge_base=None):
    module = task.get_annotation(MODULE_OBJECT_ANNOTATION_KEY)
    if task.executable_type == ExecutableType.MODULE_TYPE and module and module.arguments:

        mo = task.module_options
        module_fqcn = module.fqcn
        module_short = ""
        if module_fqcn:
            parts = module_fqcn.split(".")
            if len(parts) <= 2:
                module_short = module_fqcn.split(".")[-1]
            elif len(parts) > 2:
                module_short = ".".join(module_fqcn.split(".")[2:])
        default_args = {}
        if module_short and module_short in task.module_defaults:
            default_args = task.module_defaults[module_short]
        elif module_fqcn and module_fqcn in task.module_defaults:
            default_args = task.module_defaults[module_fqcn]
        else:
            for group_name in task.module_defaults:
                tmp_args = task.module_defaults[group_name]
                found = False
                if not group_name.startswith("group/"):
                    continue
                if knowledge_base:
                    groups = knowledge_base.kb_client.search_action_group(group_name)
                    if not groups:
                        continue
                    for group_dict in groups:
                        if not group_dict:
                            continue
                        if not isinstance(group_dict, dict):
                            continue
                        group = ActionGroupMetadata.from_dict(group_dict)
                        if module_short and module_short in group.group_modules:
                            found = True
                            default_args = tmp_args
                            break
                        elif module_fqcn and module_fqcn in group.group_modules:
                            found = True
                            default_args = tmp_args
                            break
                if found:
                    break

        used_keys = []
        if isinstance(mo, dict):
            used_keys = list(mo.keys())

        available_keys = []
        required_keys = []
        alias_reverse_map = {}
        available_args = None
        wrong_keys = []
        missing_required_keys = []
        if not is_set_fact(module_fqcn):
            if module:
                for arg in module.arguments:
                    available_keys.extend(arg.available_keys())
                    if arg.required:
                        aliases = arg.aliases if arg.aliases else []
                        req_k = {"key": arg.name, "aliases": aliases}
                        required_keys.append(req_k)
                    if arg.aliases:
                        for al in arg.aliases:
                            alias_reverse_map[al] = arg.name
                available_args = module.arguments

            wrong_keys = [key for key in used_keys if key not in available_keys]

            for k in required_keys:
                name = k.get("key", "")
                aliases = k.get("aliases", [])
                if name in used_keys:
                    continue
                if name in default_args:
                    continue
                if aliases:
                    found = False
                    for a_k in aliases:
                        if a_k in used_keys:
                            found = True
                            break
                        if a_k in default_args:
                            found = True
                            break
                    if found:
                        continue
                # here, the required key was not found
                missing_required_keys.append(name)

        used_alias_and_real_keys = []
        for k in used_keys:
            if k not in alias_reverse_map:
                continue
            real_name = alias_reverse_map[k]
            used_alias_and_real_keys.append(
                {
                    "used_alias": k,
                    "real_key": real_name,
                }
            )

        task.set_annotation("module.wrong_arg_keys", wrong_keys)
        task.set_annotation("module.available_arg_keys", available_keys)
        task.set_annotation("module.required_arg_keys", required_keys)
        task.set_annotation("module.missing_required_arg_keys", missing_required_keys)
        # task.set_annotation("module.available_args", available_args)
        # task.set_annotation("module.default_args", default_args)
        task.set_annotation("module.used_alias_and_real_keys", used_alias_and_real_keys)
    return task


# P003 rule in ARI
def set_module_arg_value_annotations(task: Task):
    module = task.get_annotation(MODULE_OBJECT_ANNOTATION_KEY)
    if task.executable_type == ExecutableType.MODULE_TYPE and module and module.arguments:

        wrong_values = []
        undefined_values = []
        unknown_type_values = []
        module_options = task.module_options
        arguments = task.get_annotation(ARGUMENTS_ANNOTATION_KEY)
        if isinstance(module_options, dict):
            for key in module_options:
                raw_value = module_options[key]
                resolved_value = None
                if arguments and len(arguments.templated) >= 1:
                    resolved_value = arguments.templated[0][key]
                spec = None
                for arg_spec in module.arguments:
                    if key == arg_spec.name or (arg_spec.aliases and key in arg_spec.aliases):
                        spec = arg_spec
                        break
                if not spec:
                    continue

                d = {"key": key}
                wrong_val = False
                unknown_type_val = False
                if spec.type:
                    actual_type = ""
                    # if the raw_value is not a variable
                    if not isinstance(raw_value, str) or "{{" not in raw_value:
                        actual_type = type(raw_value).__name__
                    else:
                        # otherwise, check the resolved value
                        # if the variable could not be resovled successfully
                        if isinstance(resolved_value, str) and "{{" in resolved_value:
                            pass
                        elif is_loop_var(raw_value, task):
                            # if the variable is loop var, use the element type as actual type
                            resolved_element = None
                            if resolved_value:
                                if isinstance(resolved_value, list):
                                    resolved_element = resolved_value[0]
                                else:
                                    resolved_element = resolved_value
                            if resolved_element:
                                actual_type = type(resolved_element).__name__
                        else:
                            # otherwise, use the resolved value type as actual type
                            actual_type = type(resolved_value).__name__

                    if actual_type:
                        type_wrong = False
                        if spec.type != "any" and actual_type != spec.type:
                            type_wrong = True
                        elements_type_wrong = False
                        no_elements = False
                        if spec.elements:
                            if spec.elements != "any" and actual_type != spec.elements:
                                elements_type_wrong = True
                        else:
                            no_elements = True
                        if type_wrong and (elements_type_wrong or no_elements):
                            d["expected_type"] = spec.type
                            d["actual_type"] = actual_type
                            d["actual_value"] = raw_value
                            wrong_val = True
                    else:
                        d["expected_type"] = spec.type
                        d["unknown_type_value"] = resolved_value
                        unknown_type_val = True

                if wrong_val:
                    wrong_values.append(d)

                if unknown_type_val:
                    unknown_type_values.append(d)

                sub_args = None
                if arguments and isinstance(arguments, dict):
                    sub_args = arguments.get(key)
                if sub_args:
                    undefined_vars = [v.name for v in sub_args.vars if v and v.type == VariableType.Unknown]
                    if undefined_vars:
                        undefined_values.append({"key": key, "value": raw_value, "undefined_variables": undefined_vars})

        task.set_annotation("module.wrong_arg_values", wrong_values)
        task.set_annotation("module.undefined_values", undefined_values)
        task.set_annotation("module.unknown_type_values", unknown_type_values)
    return task


# P004 rule in ARI
def set_variable_annotations(task: Task):
    undefined_variables = []
    unknown_name_vars = []
    unnecessary_loop = []
    task_arg_keys = []
    arguments = task.get_annotation(ARGUMENTS_ANNOTATION_KEY)
    variable_use = task.get_annotation(VARIABLES_USED_ANNOTATION_KEY)
    if arguments and arguments.type == ArgumentsType.DICT:
        task_arg_keys = list(arguments.raw.keys())
    if variable_use:
        for v_name in variable_use:
            v = variable_use[v_name]
            if v and v[-1].type == VariableType.Unknown:
                if v_name not in undefined_variables:
                    undefined_variables.append(v_name)
                if v_name not in unknown_name_vars and v_name not in task_arg_keys:
                    unknown_name_vars.append(v_name)
                if v_name not in unnecessary_loop:
                    v_str = "{{ " + v_name + " }}"
                    if not is_loop_var(v_str, task):
                        unnecessary_loop.append({"name": v_name, "suggested": v_name.replace("item.", "")})

    task.set_annotation("variable.undefined_vars", undefined_variables)
    task.set_annotation("variable.unknown_name_vars", unknown_name_vars)
    task.set_annotation("variable.unnecessary_loop_vars", unnecessary_loop)
    return task


# remove temporary annotations (to remove large object data)
def omit_object_annotations(task: Task):
    task.delete_annotation(MODULE_OBJECT_ANNOTATION_KEY)
    task.delete_annotation(ARGUMENTS_ANNOTATION_KEY)
    return task


def is_set_fact(module_fqcn):
    return module_fqcn == "ansible.builtin.set_fact"


def is_loop_var(value, task):
    # `item` and alternative loop variable (if any) should not be replaced to avoid breaking loop
    skip_variables = ["item"]
    if task.loop and isinstance(task.loop, dict):
        skip_variables.extend(list(task.loop.keys()))

    _v = value.replace(" ", "")

    for var in skip_variables:
        for _prefix in ["}}", "|", "."]:
            pattern = "{{" + var + _prefix
            if pattern in _v:
                return True
    return False
sage_scan/process/annotations.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from dataclasses import dataclass
from sage_scan.models import (
    Task,
    Module,
    SageProject,
)
from sage_scan.process.annotations import MODULE_OBJECT_ANNOTATION_KEY
from ansible_risk_insight.models import (
    ExecutableType,
    Module as ARIModule,
    ActionGroupMetadata,
    VariableType,
)
from ansible_risk_insight.risk_assessment_model import RAMClient


SAGE_KB_DATA_DIR = os.getenv("SAGE_KB_DATA_DIR", None)

# alias
if not SAGE_KB_DATA_DIR:
    SAGE_KB_DATA_DIR = os.getenv("ARI_KB_DATA_DIR", None)


@dataclass
class KnowledgeBase(object):
    kb_client: RAMClient = None

    def __post_init__(self):
        if not self.kb_client:
            self.init_kb_client()

    def init_kb_client(self):
        if SAGE_KB_DATA_DIR is None:
            raise ValueError(f"Please specify an existing SAGE KB dir by an env param:\n$ export SAGE_KB_DATA_DIR=<PATH/TO/SAGE_KB_DATA_DIR>")

        if not os.path.exists(SAGE_KB_DATA_DIR):
            raise ValueError(f"the SAGE_KB_DATA_DIR does not exist: {SAGE_KB_DATA_DIR}")

        self.kb_client = RAMClient(root_dir=SAGE_KB_DATA_DIR)
        return

    def set_kb_client_from_ram_client(self, ram_client=None):
        if ram_client:
            self.kb_client = ram_client
        return

    def resolve_task(self, task: Task, set_module_object_annotation: bool=False):
        exec_type = task.executable_type
        include_types = [ExecutableType.ROLE_TYPE, ExecutableType.TASKFILE_TYPE]

        task = self.set_module_info(task, set_module_object_annotation)
        if exec_type in include_types:
            task = self.set_include_info(task)
        
        if exec_type == ExecutableType.MODULE_TYPE:
            if task.module_info and isinstance(task.module_info, dict):
                task.resolved_name = task.module_info.get("fqcn", "")
        elif exec_type == ExecutableType.ROLE_TYPE:
            if task.include_info and isinstance(task.include_info, dict):
                task.resolved_name = task.include_info.get("fqcn", "")
        elif exec_type == ExecutableType.TASKFILE_TYPE:
            if task.include_info and isinstance(task.include_info, dict):
                task.resolved_name = task.include_info.get("key", "")

        return task
    
    def set_module_info(self, task: Task, set_module_object_annotation: bool=False):
        if not isinstance(task, Task):
            raise ValueError(f"expect a task object, but {type(task)}")

        raw_module_name = task.module
        result = self.kb_client.search_module(name=raw_module_name)
        module = None
        if result and isinstance(result, list) and isinstance(result[0], dict):
            _module = result[0].get("object", None)
            if isinstance(_module, Module):
                module = _module
            elif isinstance(_module, ARIModule):
                module = Module.from_ari_obj(_module)
        if module:
            task.module_info = {
                "collection": module.collection,
                "short_name": module.name,
                "fqcn": module.fqcn,
                "key": module.key,
            }
            if set_module_object_annotation:
                task.set_annotation(MODULE_OBJECT_ANNOTATION_KEY, module)

        return task
    
    def set_include_info(self, task: Task):
        if not isinstance(task, Task):
            raise ValueError(f"expect a task object, but {type(task)}")

        exec_target = task.executable
        exec_type = task.executable_type
        include_info = {}
        if exec_type == ExecutableType.ROLE_TYPE:
            result = self.kb_client.search_role(name=exec_target)
            if not result:
                return task
            
            if not isinstance(result, list):
                return task
            
            if not isinstance(result[0], dict):
                return task
        
            role = result[0].get("object", None)
            if not role:
                return task
            
            include_info = {
                "type": "role",
                "fqcn": role.fqcn,
                "path": role.defined_in,
                "key": role.spec.key,
            }

        elif exec_type == ExecutableType.TASKFILE_TYPE:
            result = self.kb_client.search_taskfile(name=exec_target, is_key=True)
            if not result:
                return task
            
            if not isinstance(result, list):
                return task
            
            if not isinstance(result[0], dict):
                return task
        
            taskfile = result[0].get("object", None)
            if not taskfile:
                return task
            
            include_info = {
                "type": "taskfile",
                "path": taskfile.defined_in,
                "key": taskfile.key,
            }
        
        task.include_info = include_info
        return task



sage_scan/process/knowledge_base.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from sage_scan.models import SageObject, SageProject, Playbook, TaskFile, Play, Task, Role
from sage_scan.process.variable_resolver import VariableResolver
from sage_scan.process.knowledge_base import KnowledgeBase, MODULE_OBJECT_ANNOTATION_KEY
from sage_scan.process.annotations import (
    set_module_spec_annotations,
    set_module_arg_key_annotations,
    set_module_arg_value_annotations,
    set_variable_annotations,
    omit_object_annotations,
)
from ansible_risk_insight.models import Annotation


# find all tasks defined in the specified playbook from SageProject
def get_tasks_in_playbook(playbook: Playbook, project: SageProject=None):
    play_keys = playbook.plays
    tasks = []
    for p_key in play_keys:
        play = project.get_object(key=p_key)
        if not play:
            continue
        p_tasks = get_tasks_in_play(play, project)
        tasks.extend(p_tasks)
    return tasks


# find all tasks defined in the specified play from SageProject
def get_tasks_in_play(play: Play, project: SageProject=None):
    task_keys = play.pre_tasks + play.tasks + play.post_tasks
    tasks = []
    for t_key in task_keys:
        task = project.get_object(key=t_key)
        if not task:
            continue
        tasks.append(task)
    return tasks


# find all tasks defined in the specified taskfile from SageProject
def get_tasks_in_taskfile(taskfile: TaskFile, project: SageProject=None):
    task_keys = taskfile.tasks
    tasks = []
    for t_key in task_keys:
        task = project.get_object(key=t_key)
        if not task:
            continue
        tasks.append(task)
    return tasks


# find all tasks defined in the specified playbook or taskfile from SageProject
def get_tasks_in_file(target: Playbook|TaskFile=None, project: SageProject=None):
    tasks = []
    if isinstance(target, Playbook):
        tasks = get_tasks_in_playbook(target, project)
    elif isinstance(target, TaskFile):
        tasks = get_tasks_in_taskfile(target, project)
    return tasks


# find all tasks defined in the specified playbook or taskfile from SageProject
# if `root` is an object key, get the object from SageProject first
def get_tasks(root: str | SageObject, project: SageProject=None):
    root_obj = root
    if isinstance(root, str):
        root_obj = project.get_object(key=root)
    return get_tasks_in_file(target=root_obj, project=project)


# find all plays defined in the specified playbook from SageProject
def get_plays(playbook: Playbook, project: SageProject):
    play_keys = playbook.plays
    plays = []
    for p_key in play_keys:
        play = project.get_object(key=p_key)
        if not play:
            continue
        if not isinstance(play, Play):
            continue
        plays.append(play)
    return plays


# find all taskfiles in the speciifed role from SageProject
def get_taskfiles_in_role(role: Role, project: SageProject):
    taskfile_keys = role.taskfiles
    taskfiles = []
    for tf_key in taskfile_keys:
        taskfile = project.get_object(key=tf_key)
        if not taskfile:
            continue
        if not isinstance(taskfile, TaskFile):
            continue
        taskfiles.append(taskfile)
    return taskfiles


# find main.yml or main.yaml in the specified role from SageProject
def get_main_taskfile_for_role(role: Role, project: SageProject):
    taskfiles = get_taskfiles_in_role(role, project)
    for tf in taskfiles:
        filename = os.path.basename(tf.filepath)
        if filename in ["main.yml", "main.yaml"]:
            return tf
    return None


# find a parent role for the specified taskfile if it exists
def find_parent_role(taskfile: TaskFile, project: SageProject):
    for role in project.roles:
        taskfile_keys = role.taskfiles
        if taskfile.key in taskfile_keys:
            return role
    return None


# get call tree which starts from the specified entrypoint
# call tree is a list of (obj, child_obj)
def get_call_tree_by_entrypoint(entrypoint: Playbook|Role|TaskFile, project: SageProject, follow_include: bool=True):
    return project.get_call_tree_by_entrypoint(entrypoint, follow_include)


# get all call sequences found in the SageProject
# call sequence is a sequence of objects executed by an entrypoint
# e.g.) Playbook --> Play 1 -> Task 1a -> Task 1b -> Play 2 -> Task 2a 
def get_all_call_sequences(project: SageProject, follow_include: bool=True):
    return project.get_all_call_sequences(follow_include)


# get a call sequence which contains the specified task
def get_call_sequence_for_task(task: Task, project: SageProject, follow_include: bool=True):
    return project.get_call_sequence_for_task(task, follow_include)


# get call sequence which starts from the specified entrypoint
def get_call_sequence_by_entrypoint(entrypoint: Playbook|Role|TaskFile, project: SageProject, follow_include: bool=True):
    return project.get_call_sequence_by_entrypoint(entrypoint, follow_include)


# get task sequence which starts from the specified entrypoint
def get_task_sequence_by_entrypoint(entrypoint: Playbook|Role|TaskFile, project: SageProject, follow_include: bool=True):
    call_seq = get_call_sequence_by_entrypoint(entrypoint, project, follow_include)
    if not call_seq:
        return None
    return [obj for obj in call_seq if isinstance(obj, Task)]


# get a task sequence which starts from the specified playbook
def get_task_sequence_for_playbook(playbook: Playbook, project: SageProject, follow_include: bool=True):
    return get_task_sequence_by_entrypoint(playbook, project, follow_include)


# get a task sequences which starts from the specified role
# this returns a list of task sequences; each sequence starts from a single taskfile in the role
def get_task_sequences_for_role(role: Role, project: SageProject, follow_include: bool=True):
    taskfiles = get_taskfiles_in_role(role, project)
    task_seq_list = []
    for taskfile in taskfiles:
        task_seq = get_task_sequence_by_entrypoint(taskfile, project, follow_include)
        if not task_seq:
            continue
        task_seq_list.append(task_seq)
    return task_seq_list


# get a task sequences which starts from the specified taskfile
def get_task_sequence_for_taskfile(taskfile: TaskFile, project: SageProject, follow_include: bool=True):
    return get_task_sequence_by_entrypoint(taskfile, project, follow_include)


# embed `defined_vars` annotation to all objects in a call sequence
def set_defined_vars(call_seq: list):
    resolver = VariableResolver()
    obj_and_vars_list = resolver.set_defined_vars(call_seq=call_seq)
    call_seq = resolver.call_seq
    return obj_and_vars_list


# get defined vars for the specified object
# if object is a TaskFile, this returnes defined vars of a Role which contains the TaskFile if found
def get_defined_vars(object: SageObject, project: SageProject):
    target = object
    obj_list = [object]
    if isinstance(object, TaskFile):
        role = find_parent_role(object, project)
        if role:
            target = role
            obj_list = [role, object]

    resolver = VariableResolver()
    return resolver.get_defined_vars(object=target, call_seq=obj_list)

# get used vars for the specified object
# this returns a dict of var_name and var_value pairs, where values are resolved as much as possible
# if variable resolution fails, the value will be a placeholder string like `{{ var_name }}`
def get_used_vars(object: SageObject, project: SageProject, follow_include: bool=True):
    obj_list = []
    if isinstance(object, TaskFile):
        role = find_parent_role(object, project)
        if role:
            obj_list = [role]
    call_seq = get_call_sequence_by_entrypoint(
        entrypoint=object,
        project=project,
        follow_include=follow_include,
    )
    obj_list.extend(call_seq)
    if not obj_list:
        return {}
    
    target = obj_list[-1]
    resolver = VariableResolver()
    return resolver.get_used_vars(object=target, call_seq=obj_list)


def set_vars_annotation(object: SageObject, project: SageProject):
    obj_list = []
    if isinstance(object, TaskFile):
        role = find_parent_role(object, project)
        if role:
            obj_list = [role]
    call_seq = get_call_sequence_by_entrypoint(entrypoint=object, project=project)
    obj_list.extend(call_seq)
    if not obj_list:
        return []
    
    resolver = VariableResolver()
    obj_vars_list = resolver.set_used_vars(call_seq=obj_list)
    obj_list = [obj for (obj, _, _) in obj_vars_list]
    return obj_list


# returns all entrypoint objects
# playbooks, roles and independent taskfiles (=not in a role) can be an entrypoint
def list_entrypoints(project: SageProject):
    entrypoints = []
    entrypoints.extend(project.playbooks)
    entrypoints.extend(project.roles)
    # only independent taskfiles; skip taskfiles in role
    entrypoints.extend([tf for tf in project.taskfiles if not tf.role])
    return entrypoints


# set `module_info` in task
# this requires ARI KB data for non-builtin modules
def set_module_info_to_task(task: Task):
    kb = KnowledgeBase()
    return kb.set_module_info(task)


def set_primary_annotations_to_project(project: SageProject):
    resolver = VariableResolver()
    # set variable data
    project = resolver.resolve_all_vars_in_project(project=project)
    kb = KnowledgeBase()
    tasks = project.tasks
    for task in tasks:
        # set module_info/include_info
        kb.resolve_task(task, set_module_object_annotation=True)

        # set P001 annotations
        set_module_spec_annotations(task)

        # set P002 annotations
        set_module_arg_key_annotations(task, knowledge_base=kb)

        # set P003 annotations
        set_module_arg_value_annotations(task)

        # set P004 annotations
        set_variable_annotations(task)

        # remove temporary annotations (to avoid saving large data)
        omit_object_annotations(task)
    return project

sage_scan/process/utils.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import re
import copy
from pathlib import Path
from dataclasses import dataclass, field
from sage_scan.models import (
    Collection,
    Playbook,
    Play,
    TaskFile,
    Role,
    Task,
    SageObject,
    SageProject,
)
from sage_scan.process.annotations import (
    ARGUMENTS_ANNOTATION_KEY,
    VARIABLES_SET_ANNOTATION_KEY,
    VARIABLES_USED_ANNOTATION_KEY,
    is_loop_var,
)
from ansible_risk_insight.models import (
    Variable,
    VariableType,
    BecomeInfo,
    InventoryType,
    Arguments,
    ArgumentsType,
    immutable_var_types,
)


variable_block_re = re.compile(r"{{[^}]+}}")
p = Path(__file__).resolve().parent
ansible_special_variables = [line.replace("\n", "") for line in open(p / "ansible_variables.txt", "r").read().splitlines()]
variable_block_re = re.compile(r"{{[^}]+}}")


@dataclass
class VariableResolver(object):

    def resolve_all_vars_in_project(self, project: SageProject):
        all_call_sequences = project.get_all_call_sequences()
        for call_seq in all_call_sequences:
            self.traverse(call_seq=call_seq)
        return project

    def get_defined_vars(self, object: SageObject, call_seq: list):
        obj_and_vars_list = self.set_defined_vars(call_seq=call_seq)
        for obj, defnied_vars, _ in obj_and_vars_list:
            if obj.key == object.key:
                return defnied_vars
        return {}
    
    def get_used_vars(self, object: SageObject, call_seq: list):
        obj_and_vars_list = self.set_used_vars(call_seq=call_seq)
        for obj, _, used_vars in obj_and_vars_list:
            if obj.key == object.key:
                return used_vars
        return {}

    def set_defined_vars(self, call_seq: list):
        obj_and_vars_list = self.traverse(call_seq=call_seq)
        return obj_and_vars_list
    
    def set_used_vars(self, call_seq: list):
        obj_and_vars_list = self.traverse(call_seq=call_seq)
        return obj_and_vars_list

    def traverse(self, call_seq: list):
        context = VariableContext()
        obj_and_vars_list = []
        for obj in call_seq:
            context.add(obj)
            if isinstance(obj, Task):
                task = obj
                resolved = resolve_module_options(context, task)
                resolved_module_options = resolved[0]
                resolved_variables = resolved[1]
                used_variables = resolved[3]

                _vars = []
                is_mutable = False
                for rv in resolved_variables:
                    v_name = rv.get("key", "")
                    v_value = rv.get("value", "")
                    v_type = rv.get("type", VariableType.Unknown)
                    elements = []
                    if v_name in used_variables:
                        if not isinstance(used_variables[v_name], dict):
                            continue
                        for u_v_name, info in used_variables[v_name].items():
                            if u_v_name == v_name:
                                continue
                            u_v_value = info.get("value", "")
                            u_v_type = info.get("type", VariableType.Unknown)
                            u_v = Variable(
                                name=u_v_name,
                                value=u_v_value,
                                type=u_v_type,
                                used_in=task.key,
                            )
                            elements.append(u_v)
                    v = Variable(
                        name=v_name,
                        value=v_value,
                        type=v_type,
                        elements=elements,
                        used_in=task.key,
                    )
                    _vars.append(v)
                    if v.is_mutable:
                        is_mutable = True

                for v in _vars:
                    history = context.var_use_history.get(v.name, [])
                    history.append(v)
                    context.var_use_history[v.name] = history

                m_opts = task.module_options
                if isinstance(m_opts, list):
                    args_type = ArgumentsType.LIST
                elif isinstance(m_opts, dict):
                    args_type = ArgumentsType.DICT
                else:
                    args_type = ArgumentsType.SIMPLE
                arguments = Arguments(
                    type=args_type,
                    raw=m_opts,
                    vars=_vars,
                    resolved=True,  # TODO: False if not resolved
                    templated=resolved_module_options,
                    is_mutable=is_mutable,
                )
                # deep copy the history here because the context is updated by subsequent taskcalls
                defined_vars = copy.deepcopy(context.var_set_history)
                used_vars = copy.deepcopy(context.var_use_history)
                task.set_annotation(VARIABLES_SET_ANNOTATION_KEY, defined_vars)
                task.set_annotation(VARIABLES_USED_ANNOTATION_KEY, used_vars)
                task.set_annotation(ARGUMENTS_ANNOTATION_KEY, arguments)

                defined_vars_key_value = {k: v[0].value for k, v in defined_vars.items() if v and isinstance(v[0], Variable)}
                used_vars_key_value = {}
                for k, v in used_vars.items():
                    if not v:
                        continue
                    _var = v[0]
                    if not isinstance(_var, Variable):
                        continue

                    # skip if this var is registered one
                    parent_var_name = _var.name
                    if "." in parent_var_name:
                        parent_var_name = parent_var_name.split(".")[0]
                    is_registered_var = False
                    for p, q in defined_vars.items():
                        if not q:
                            continue
                        _var_q = q[0]
                        if not isinstance(_var_q, Variable):
                            continue
                        if _var_q.name == parent_var_name and _var_q.type == VariableType.RegisteredVars:
                            is_registered_var = True
                            break
                    if is_registered_var:
                        continue

                    # skip loop var by type
                    if _var.type == VariableType.LoopVars:
                        continue

                    # skip loop var by name
                    # (this may be removed in the future because the type check above could be enough)
                    var_block = "{{ " + k + " }}"
                    if is_loop_var(var_block, task):
                        continue
                    
                    value = _var.value
                    if _var.type == VariableType.Unknown:
                        value = make_value_placeholder(k)
                    used_vars_key_value[k] = value
                    
                obj_and_vars_list.append((obj, defined_vars_key_value, used_vars_key_value))
            else:
                last_defined = {}
                last_used = {}
                if obj_and_vars_list:
                    _, last_defined, last_used = obj_and_vars_list[-1]
                obj_and_vars_list.append((obj, last_defined, last_used))
        return obj_and_vars_list
    
    def update_resolved_vars_dict(self, vars_dict, var_name, var_value):
        def _recursive_update(d, keys, value):
            if not isinstance(d, dict):
                d = {}
            new_d = d.copy()

            if not keys:
                return {}
            key = keys[0]
            if len(keys) <= 1:
                new_d[key] = value
            else:
                rest_keys = keys[1:]
                if key not in new_d or not isinstance(new_d[key], dict):
                    new_d[key] = {}
                new_d[key] = _recursive_update(new_d[key], rest_keys, value)
            return new_d

        updated = vars_dict.copy()
        if "." in var_name:
            parts = var_name.split(".")
            updated = _recursive_update(updated, parts, var_value)
        else:
            updated[var_name] = var_value
        return updated


def make_value_placeholder(var_name: str):
    if "." in var_name:
        var_name = var_name.replace(".", "_")
    return "{{ " + var_name + " }}"
    

def flatten_vars_dict(vars_dict: dict, _prefix: str = ""):
    flat_vars_dict = {}
    for k, v in vars_dict.items():
        if isinstance(v, dict):
            flat_var_name = f"{_prefix}{k}"
            flat_vars_dict.update({flat_var_name: v})
            new_prefix = f"{flat_var_name}."
            _tmp = flatten_vars_dict(v, new_prefix)
            flat_vars_dict.update(_tmp)
        else:
            flat_key = f"{_prefix}{k}"
            flat_vars_dict.update({flat_key: v})
    return flat_vars_dict


def extract_variable_names(txt):
    if not variable_block_re.search(txt):
        return []
    found_var_blocks = variable_block_re.findall(txt)
    blocks = []
    for b in found_var_blocks:
        if "lookup(" in b.replace(" ", ""):
            continue
        parts = b.split("|")
        var_name = ""
        default_var_name = ""
        for i, p in enumerate(parts):
            if i == 0:
                var_name = p.replace("{{", "").replace("}}", "")
                if " if " in var_name and " else " in var_name:
                    # this block is not just a variable, but an expression
                    # we need to split this with a space to get its elements
                    skip_elements = ["if", "else", "+", "is", "defined"]
                    sub_parts = var_name.split(" ")
                    for sp in sub_parts:
                        if not sp:
                            continue
                        if sp and sp in skip_elements:
                            continue
                        if sp and sp[0] in ['"', "'"]:
                            continue
                        var_name = sp
                        break
                var_name = var_name.replace(" ", "")
                if "lookup(" in var_name and "first_found" in var_name:
                    var_name = var_name.split(",")[-1].replace(")", "")
                if var_name and var_name[0] == "(":
                    var_name = var_name.split(")")[0].replace("(", "")
                if "+" in var_name:
                    sub_parts = var_name.split("+")
                    for sp in sub_parts:
                        if not sp:
                            continue
                        if sp and sp[0] in ['"', "'"]:
                            continue
                        var_name = sp
                        break
            else:
                if "default(" in p and ")" in p:
                    default_var = p.replace("}}", "").replace("default(", "").replace(")", "").replace(" ", "")
                    if not default_var.startswith('"') and not default_var.startswith("'") and not re.compile(r"[0-9].*").match(default_var):
                        default_var_name = default_var
        tmp_b = {
            "original": b,
        }
        if var_name == "":
            continue
        tmp_b["name"] = var_name
        if default_var_name != "":
            tmp_b["default"] = default_var_name
        blocks.append(tmp_b)
    return blocks


def resolved_vars_contains(resolved_vars, new_var):
    if not isinstance(new_var, dict):
        return False
    new_var_key = new_var.get("key", "")
    if new_var_key == "":
        return False
    if not isinstance(resolved_vars, list):
        return False
    for var in resolved_vars:
        if not isinstance(var, dict):
            continue
        var_key = var.get("key", "")
        if var_key == "":
            continue
        if var_key == new_var_key:
            return True
    return False


@dataclass
class VariableContext:
    keep_obj: bool = False
    variables: dict = field(default_factory=dict)
    options: dict = field(default_factory=dict)
    inventories: list = field(default_factory=list)
    role_defaults: list = field(default_factory=list)
    role_vars: list = field(default_factory=list)
    registered_vars: list = field(default_factory=list)
    set_facts: list = field(default_factory=list)
    task_vars: list = field(default_factory=list)

    become: BecomeInfo = None
    module_defaults: dict = field(default_factory=dict)

    var_set_history: dict = field(default_factory=dict)
    var_use_history: dict = field(default_factory=dict)

    _flat_vars: dict = field(default_factory=dict)

    def add(self, obj):
        _spec = None
        if isinstance(obj, SageObject):
            _spec = obj
        # variables
        if isinstance(_spec, Playbook):
            self.variables.update(_spec.variables)
            self.update_flat_vars(_spec.variables)
            for key, val in _spec.variables.items():
                current = self.var_set_history.get(key, [])
                current.append(Variable(name=key, value=val, type=VariableType.PlaybookGroupVarsAll, setter=_spec.key))
                self.var_set_history[key] = current
        elif isinstance(_spec, Play):
            self.variables.update(_spec.variables)
            self.update_flat_vars(_spec.variables)
            for key, val in _spec.variables.items():
                current = self.var_set_history.get(key, [])
                current.append(Variable(name=key, value=val, type=VariableType.PlayVars, setter=_spec.key))
                self.var_set_history[key] = current
            if _spec.become:
                self.become = _spec.become
            if _spec.module_defaults:
                self.module_defaults = _spec.module_defaults
        elif isinstance(_spec, Role):
            self.variables.update(_spec.default_variables)
            self.update_flat_vars(_spec.default_variables)
            self.variables.update(_spec.variables)
            self.update_flat_vars(_spec.variables)
            for var_name in _spec.default_variables:
                self.role_defaults.append(var_name)
            for var_name in _spec.variables:
                self.role_vars.append(var_name)
            for key, val in _spec.default_variables.items():
                current = self.var_set_history.get(key, [])
                current.append(Variable(name=key, value=val, type=VariableType.RoleDefaults, setter=_spec.key))
                self.var_set_history[key] = current
            for key, val in _spec.variables.items():
                current = self.var_set_history.get(key, [])
                current.append(Variable(name=key, value=val, type=VariableType.RoleVars, setter=_spec.key))
                self.var_set_history[key] = current
        elif isinstance(_spec, Collection):
            self.variables.update(_spec.variables)
            self.update_flat_vars(_spec.variables)
        elif isinstance(_spec, TaskFile):
            self.variables.update(_spec.variables)
            self.update_flat_vars(_spec.variables)
        elif isinstance(_spec, Task):
            self.variables.update(_spec.variables)
            self.update_flat_vars(_spec.variables)
            self.variables.update(_spec.registered_variables)
            self.update_flat_vars(_spec.registered_variables)
            self.variables.update(_spec.set_facts)
            self.update_flat_vars(_spec.set_facts)
            for var_name in _spec.registered_variables:
                self.registered_vars.append(var_name)
            for var_name in _spec.set_facts:
                self.set_facts.append(var_name)
            for key, val in _spec.variables.items():
                current = self.var_set_history.get(key, [])
                current.append(Variable(name=key, value=val, type=VariableType.TaskVars, setter=_spec.key))
                self.var_set_history[key] = current
            for key, val in _spec.registered_variables.items():
                current = self.var_set_history.get(key, [])
                current.append(Variable(name=key, value=val, type=VariableType.RegisteredVars, setter=_spec.key))
                self.var_set_history[key] = current
            for key, val in _spec.set_facts.items():
                current = self.var_set_history.get(key, [])
                current.append(Variable(name=key, value=val, type=VariableType.SetFacts, setter=_spec.key))
                self.var_set_history[key] = current
            if _spec.become:
                self.become = _spec.become
            if _spec.module_defaults:
                self.module_defaults = _spec.module_defaults
        else:
            # Module
            return
        self.options.update(_spec.options)

    def resolve_variable(self, var_name, resolve_history={}):
        if var_name in resolve_history:
            val = resolve_history[var_name].get("value", None)
            v_type = resolve_history[var_name].get("type", VariableType.Unknown)
            return val, v_type, resolve_history

        _resolve_history = resolve_history.copy()

        v_type = None
        if var_name in ansible_special_variables:
            v_type = VariableType.HostFacts
            return None, v_type, resolve_history

        if var_name in self.role_vars:
            v_type = VariableType.RoleVars
        elif var_name in self.role_defaults:
            v_type = VariableType.RoleDefaults
        elif var_name in self.registered_vars:
            v_type = VariableType.RegisteredVars
        elif var_name in self.set_facts:
            v_type = VariableType.SetFacts
        else:
            v_type = VariableType.TaskVars

        val = self.variables.get(var_name, None)
        if val is not None:
            _resolve_history[var_name] = {"value": val, "type": v_type}

            if isinstance(val, str):
                resolved_val, _resolve_history = self.resolve_single_variable(val, _resolve_history)
                return resolved_val, v_type, _resolve_history
            elif isinstance(val, list):
                resolved_val_list = []
                for vi in val:
                    resolved_val, _resolve_history = self.resolve_single_variable(vi, _resolve_history)
                    resolved_val_list.append(resolved_val)
                return resolved_val_list, v_type, _resolve_history
            else:
                return val, v_type, _resolve_history

        val = self._flat_vars.get(var_name, None)
        if val is not None:
            _resolve_history[var_name] = {"value": val, "type": v_type}

            if isinstance(val, str):
                resolved_val, _resolve_history = self.resolve_single_variable(val, _resolve_history)
                return resolved_val, v_type, _resolve_history
            elif isinstance(val, list):
                resolved_val_list = []
                for vi in val:
                    resolved_val, _resolve_history = self.resolve_single_variable(vi, _resolve_history)
                    resolved_val_list.append(resolved_val)
                return resolved_val_list, v_type, _resolve_history
            else:
                return val, v_type, _resolve_history

        # TODO: consider group
        inventory_for_all = [iv for iv in self.inventories if iv.inventory_type == InventoryType.GROUP_VARS_TYPE and iv.name == "all"]
        for iv in inventory_for_all:
            iv_var_dict = flatten_vars_dict(iv.variables)
            val = iv_var_dict.get(var_name, None)

            if val is not None:
                _resolve_history[var_name] = {"value": val, "type": v_type}
                v_type = VariableType.InventoryGroupVarsAll
                if isinstance(val, str):
                    resolved_val, _resolve_history = self.resolve_single_variable(val, _resolve_history)
                    return resolved_val, v_type, _resolve_history
                elif isinstance(val, list):
                    resolved_val_list = []
                    for vi in val:
                        resolved_val, _resolve_history = self.resolve_single_variable(vi, _resolve_history)
                        resolved_val_list.append(resolved_val)
                    return resolved_val_list, v_type, _resolve_history
                else:
                    return val, v_type, _resolve_history

        _resolve_history[var_name] = {"value": None, "type": VariableType.Unknown}

        return None, VariableType.Unknown, _resolve_history

    def resolve_single_variable(self, txt, resolve_history=[]):
        new_history = resolve_history.copy()
        if not isinstance(txt, str):
            return txt, new_history
        if "{{" in txt:
            var_names_in_txt = extract_variable_names(txt)
            if len(var_names_in_txt) == 0:
                return txt, new_history
            resolved_txt = txt
            for var_name_in_txt in var_names_in_txt:
                original_block = var_name_in_txt.get("original", "")
                var_name = var_name_in_txt.get("name", "")
                default_var_name = var_name_in_txt.get("default", "")
                var_val_in_txt, _, new_history = self.resolve_variable(var_name, new_history)
                if var_val_in_txt is None and default_var_name != "":
                    var_val_in_txt, _, new_history = self.resolve_variable(default_var_name, new_history)
                if var_val_in_txt is None:
                    return resolved_txt, new_history
                if txt == original_block:
                    return var_val_in_txt, new_history
                resolved_txt = resolved_txt.replace(original_block, str(var_val_in_txt))
            return resolved_txt, new_history
        else:
            return txt, new_history

    def update_flat_vars(self, new_vars: dict, _prefix: str = ""):
        for k, v in new_vars.items():
            if isinstance(v, dict):
                flat_var_name = f"{_prefix}{k}"
                self._flat_vars.update({flat_var_name: v})
                new_prefix = f"{flat_var_name}."
                self.update_flat_vars(v, new_prefix)
            else:
                flat_key = f"{_prefix}{k}"
                self._flat_vars.update({flat_key: v})
        return

    def copy(self):
        return VariableContext(
            keep_obj=self.keep_obj,
            variables=copy.copy(self.variables),
            options=copy.copy(self.options),
            inventories=copy.copy(self.inventories),
            role_defaults=copy.copy(self.role_defaults),
            role_vars=copy.copy(self.role_vars),
            registered_vars=copy.copy(self.registered_vars),
        )
        # return copy.deepcopy(self)



def resolve_module_options(context: VariableContext, task: Task):
    resolved_vars = []
    variables_in_loop = []
    used_variables = {}
    if len(task.loop) == 0:
        variables_in_loop = [{}]
    else:
        loop_key = list(task.loop.keys())[0]
        loop_values = task.loop.get(loop_key, [])
        new_var = {
            "key": loop_key,
            "value": loop_values,
            "type": VariableType.LoopVars,
        }
        if not resolved_vars_contains(resolved_vars, new_var):
            resolved_vars.append(new_var)
        if isinstance(loop_values, str):
            var_names = extract_variable_names(loop_values)
            if len(var_names) == 0:
                variables_in_loop.append({loop_key: loop_values})
            else:
                var_name = var_names[0].get("name", "")
                resolved_vars_in_item, v_type, resolve_history = context.resolve_variable(var_name)
                used_variables[var_name] = resolve_history
                new_var = {
                    "key": var_name,
                    "value": resolved_vars_in_item,
                    "type": v_type,
                }
                if not resolved_vars_contains(resolved_vars, new_var):
                    resolved_vars.append(new_var)
                if isinstance(resolved_vars_in_item, list):
                    for vi in resolved_vars_in_item:
                        variables_in_loop.append(
                            {
                                loop_key: vi,
                                "__v_type__": v_type,
                                "__v_name__": var_name,
                            }
                        )
                if isinstance(resolved_vars_in_item, dict):
                    for vi_key, vi_value in resolved_vars_in_item.items():
                        variables_in_loop.append(
                            {
                                loop_key + ".key": vi_key,
                                loop_key + ".value": vi_value,
                                "__v_type__": v_type,
                            }
                        )
                else:
                    variables_in_loop.append(
                        {
                            loop_key: resolved_vars_in_item,
                            "__v_type__": v_type,
                            "__v_name__": var_name,
                        }
                    )
        elif isinstance(loop_values, list):
            for v in loop_values:
                if isinstance(v, str) and variable_block_re.search(v):
                    var_names = extract_variable_names(v)
                    if len(var_names) == 0:
                        variables_in_loop.append({loop_key: v})
                        continue
                    var_name = var_names[0].get("name", "")
                    resolved_vars_in_item, v_type, resolve_history = context.resolve_variable(var_name)
                    used_variables[var_name] = resolve_history
                    new_var = {
                        "key": var_name,
                        "value": resolved_vars_in_item,
                        "type": v_type,
                    }
                    if not resolved_vars_contains(resolved_vars, new_var):
                        resolved_vars.append(new_var)
                    if not isinstance(resolved_vars_in_item, list):
                        variables_in_loop.append(
                            {
                                loop_key: resolved_vars_in_item,
                                "__v_type__": v_type,
                                "__v_name__": var_name,
                            }
                        )
                        continue
                    for vi in resolved_vars_in_item:
                        variables_in_loop.append(
                            {
                                loop_key: vi,
                                "__v_type__": v_type,
                                "__v_name__": var_name,
                            }
                        )
                else:
                    if isinstance(v, dict):
                        tmp_variables = {}
                        for k2, v2 in v.items():
                            key = "{}.{}".format(loop_key, k2)
                            tmp_variables.update({key: v2})
                        variables_in_loop.append(tmp_variables)
                    else:
                        variables_in_loop.append({loop_key: v})
        elif isinstance(loop_values, dict):
            tmp_variables = {}
            for k, v in loop_values.items():
                key = "{}.{}".format(loop_key, k)
                tmp_variables.update({key: v})
            variables_in_loop.append(tmp_variables)
        else:
            if loop_values:
                raise ValueError("loop_values of type {} is not supported yet".format(type(loop_values).__name__))

    resolved_opts_in_loop = []
    mutable_vars_per_mo = {}
    for variables in variables_in_loop:
        resolved_opts = None
        if isinstance(task.module_options, dict):
            resolved_opts = {}
            for (
                module_opt_key,
                module_opt_val,
            ) in task.module_options.items():
                if not isinstance(module_opt_val, str):
                    resolved_opts[module_opt_key] = module_opt_val
                    continue
                if not variable_block_re.search(module_opt_val):
                    resolved_opts[module_opt_key] = module_opt_val
                    continue
                # if variables are used in the module option value string
                var_names = extract_variable_names(module_opt_val)
                resolved_opt_val = module_opt_val
                for var_name_dict in var_names:
                    original_block = var_name_dict.get("original", "")
                    var_name = var_name_dict.get("name", "")
                    default_var_name = var_name_dict.get("default", "")
                    resolved_var_val = variables.get(var_name, None)
                    if resolved_var_val is not None:
                        loop_var_type = variables.get("__v_type__", VariableType.Unknown)
                        loop_var_name = variables.get("__v_name__", "")
                        if loop_var_type not in immutable_var_types:
                            if module_opt_key not in mutable_vars_per_mo:
                                mutable_vars_per_mo[module_opt_key] = []
                            mutable_vars_per_mo[module_opt_key].append(loop_var_name)
                    if resolved_var_val is None:
                        resolved_var_val, v_type, resolve_history = context.resolve_variable(var_name)
                        used_variables[var_name] = resolve_history
                        if resolved_var_val is not None:
                            new_var = {
                                "key": var_name,
                                "value": resolved_var_val,
                                "type": v_type,
                            }
                            if not resolved_vars_contains(resolved_vars, new_var):
                                resolved_vars.append(new_var)
                            if v_type not in immutable_var_types:
                                if module_opt_key not in mutable_vars_per_mo:
                                    mutable_vars_per_mo[module_opt_key] = []
                                mutable_vars_per_mo[module_opt_key].append(var_name)
                    if resolved_var_val is None and default_var_name != "":
                        resolved_var_val, v_type, resolve_history = context.resolve_variable(default_var_name)
                        used_variables[default_var_name] = resolve_history
                        if resolved_var_val is not None:
                            new_var = {
                                "key": default_var_name,
                                "value": resolved_var_val,
                                "type": v_type,
                            }
                            if not resolved_vars_contains(resolved_vars, new_var):
                                resolved_vars.append(new_var)
                            if v_type not in immutable_var_types:
                                if module_opt_key not in mutable_vars_per_mo:
                                    mutable_vars_per_mo[module_opt_key] = []
                                mutable_vars_per_mo[module_opt_key].append(var_name)
                    if resolved_var_val is None:
                        new_var = {
                            "key": var_name,
                            "value": None,
                            "type": v_type,
                        }
                        if not resolved_vars_contains(resolved_vars, new_var):
                            resolved_vars.append(new_var)
                        continue
                    if resolved_opt_val == original_block:
                        resolved_opt_val = resolved_var_val
                        break
                    resolved_opt_val = resolved_opt_val.replace(original_block, str(resolved_var_val))
                resolved_opts[module_opt_key] = resolved_opt_val
        elif isinstance(task.module_options, str):
            resolved_opt_val = task.module_options
            if variable_block_re.search(resolved_opt_val):
                var_names = extract_variable_names(task.module_options)
                for var_name_dict in var_names:
                    original_block = var_name_dict.get("original", "")
                    var_name = var_name_dict.get("name", "")
                    default_var_name = var_name_dict.get("default", "")
                    resolved_var_val = variables.get(var_name, None)
                    if resolved_var_val is not None:
                        loop_var_type = variables.get("__v_type__", VariableType.Unknown)
                        loop_var_name = variables.get("__v_name__", "")
                        if loop_var_type not in immutable_var_types:
                            if "" not in mutable_vars_per_mo:
                                mutable_vars_per_mo[""] = []
                            mutable_vars_per_mo[""].append(loop_var_name)
                    if resolved_var_val is None:
                        resolved_var_val, v_type, resolve_history = context.resolve_variable(var_name)
                        used_variables[var_name] = resolve_history
                        if resolved_var_val is not None:
                            new_var = {
                                "key": var_name,
                                "value": resolved_var_val,
                                "type": v_type,
                            }
                            if not resolved_vars_contains(resolved_vars, new_var):
                                resolved_vars.append(new_var)
                            if v_type not in immutable_var_types:
                                if "" not in mutable_vars_per_mo:
                                    mutable_vars_per_mo[""] = []
                                mutable_vars_per_mo[""].append(var_name)
                    if resolved_var_val is None and default_var_name != "":
                        resolved_var_val, v_type, resolve_history = context.resolve_variable(default_var_name)
                        used_variables[default_var_name] = resolve_history
                        if resolved_var_val is not None:
                            new_var = {
                                "key": default_var_name,
                                "value": resolved_var_val,
                                "type": v_type,
                            }
                            if not resolved_vars_contains(resolved_vars, new_var):
                                resolved_vars.append(new_var)
                            if v_type not in immutable_var_types:
                                if "" not in mutable_vars_per_mo:
                                    mutable_vars_per_mo[""] = []
                                mutable_vars_per_mo[""].append(var_name)
                    if resolved_var_val is None:
                        new_var = {
                            "key": var_name,
                            "value": None,
                            "type": v_type,
                        }
                        if not resolved_vars_contains(resolved_vars, new_var):
                            resolved_vars.append(new_var)
                        continue
                    if resolved_opt_val == original_block:
                        resolved_opt_val = resolved_var_val
                        break
                    resolved_opt_val = resolved_opt_val.replace(original_block, str(resolved_var_val))
            resolved_opts = resolved_opt_val
        else:
            resolved_opts = task.module_options
        resolved_opts_in_loop.append(resolved_opts)
    return resolved_opts_in_loop, resolved_vars, mutable_vars_per_mo, used_variables
sage_scan/process/variable_resolver.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import tarfile
import json

# list target file names from tar.gz
def get_target_files_from_gzip(tar_gz_file, target_filename):
    files = []
    with tarfile.open(tar_gz_file, 'r:gz') as tar:
        for member in tar.getmembers():
            if member.name.endswith(target_filename):
                files.append(member.name)
    return files


# load target file content from tar.gz
def load_file_contents(member, tar_file:tarfile.TarFile):
    loaded_contents = []
    # read file in binary mode
    file_contents = tar_file.extractfile(member).readlines()
    for file_content in file_contents:
        file_content_str = file_content.decode('utf-8')
        try:
            content_dict = json.loads(file_content_str)
            loaded_contents.append(content_dict)
        except json.JSONDecodeError as e:
            print(f"JSONDecodeError: {e}")  
    return loaded_contents
sage_scan/tools/file_utils.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from mdutils.mdutils import MdUtils
from pathlib import Path
import os
import json
import argparse
from dataclasses import dataclass, field
import glob
import tempfile
import os
import collections
import jsonpickle
from report_models import ScanReport, ProjectSource, TaskCount, FileCount, \
    OtherCount, ScanCount, ErrorCount, RoleCount, FileResult, StateCount
import tarfile
from file_utils import get_target_files_from_gzip, load_file_contents

OBJ_FILE="sage-objects.json"
META_FILE="sage-metadata.json"

class SKIP_REASON:
    NO_TASK = "no task included"
    NO_TASK_desc = "no task included in the file"
    INVALID_TASKFILE = "invalid taskfile"
    INVALID_PLAYBOOK = "invalid playbook"
    INVALID_desc = "no Ansible task included and the file not in a valid path in role"
    EXT_DEP = "external dependency"
    EXT_DEP_desc = "file is in external dependency role"
    OTHERS = "other file"
    OTHERS_desc = "other type file"    
    MANY_TASKS = "too many tasks"
    MANY_TASKS_desc = "too many tasks included"
    YAMLParseError = "invalid yml"
    YAMLParseError_desc = "unable to parse YAML"
    UNKNOWN = "unknown"

class Data_Splitter:
    def __init__(self, out_dir, filename) -> None:
        self.out_dir = out_dir
        self._buffer = {}
        self.max_buffer = 100
        self.filename = filename

    def _get_lines(self, src_type, repo_name, init=False):
        if src_type is None or repo_name is None:
            return None

        v1 = None
        if src_type not in self._buffer:
            if init:
                v1 = {}
                self._buffer[src_type] = v1
            else:
                return None
        else:
            v1 = self._buffer[src_type]

        v2 = None
        if repo_name not in v1:
            if init:
                v2 = []
                v1[repo_name] = v2
            else:
                return None
        else:
            v2 = v1[repo_name]

        return v2

    def save_to_buffer(self, src_type=None, repo_name=None, str=None):
        if src_type is None or repo_name is None or str is None:
            return

        lines = self._get_lines(src_type, repo_name, init=True)
        lines.append(str)

        if len(lines) > self.max_buffer:
            self._dump_buffer(src_type, repo_name)

    def _dump_buffer(self, src_type, repo_name):
        os.makedirs(os.path.join(self.out_dir, src_type, repo_name), exist_ok=True)
        file = os.path.join(self.out_dir, src_type, repo_name, self.filename)
        lines = self._get_lines(src_type, repo_name)
        with open(file, mode="a") as f:
            for line in lines:
                f.write(f"{line.rstrip()}\n")
        self._buffer[src_type][repo_name] = []

    def save_all(self):
        for src_type in self._buffer:
            for repo_name in self._buffer[src_type]:
                self._dump_buffer(src_type, repo_name)

    def get_repo_names(self, src_type):
        return self._buffer[src_type].keys()


class ScanResultSummarizer:
    def __init__(self, outdir, json_dir, detail_repo_dir, input_dir) -> None:
        self.outdir = outdir
        self.json_outdir = json_dir
        self.md_outdir = detail_repo_dir
        self.objects_root_dir = input_dir
        self._report_buffer = []
        self.max_buffer = 100
        self.tar_file = None

    def generate_repo_summary(self, obj_file, gzip_input):
        meta_file = obj_file.replace(OBJ_FILE, META_FILE)
        summary = self.compute_scan_report(obj_file, meta_file, gzip_input)
        return summary

    def generate_repo_summary_all(self, files, repo_scan_results_path):
        gzip_input = False
        if self.tar_file:
            gzip_input = True
        for obj_file in files:
            print(f"computing repo scan result ... {obj_file}")
            repo_scan_report = self.generate_repo_summary(obj_file, gzip_input)
            # generate repo md file if not passed
            self.generate_repo_report(repo_scan_report)
            self._report_buffer.append(repo_scan_report)
            self.write_data(repo_scan_results_path)
        self.write_all(repo_scan_results_path)
        return

    def compute_scan_report(self, object_json, meta_json, gzip_input):
        if gzip_input:
            metadata = load_file_contents(meta_json, self.tar_file)
            objects = load_file_contents(object_json, self.tar_file)
        else:
            metadata = load_json_data(meta_json)
            objects =load_json_data(object_json)
            # convert to relative path
            object_json = os.path.relpath(object_json, self.objects_root_dir)
            meta_json = os.path.relpath(meta_json, self.objects_root_dir)

        if len(objects) == 0:
            print(f"{object_json} is empty")
            ps = ProjectSource(metadata[0]["source"]["type"], metadata[0]["source"]["repo_name"], object_json, meta_json)
            error = {"EmptyObjectsJsonError": 1}
            state = StateCount()
            state.unknown = 1
            return ScanReport(project_count=1, projects=[ps], error=error, state_count=state)
        # summarize yml inventory
        target_files = []
        out_scopes = []
        roles = []
        yml_files = metadata[0].get("yml_files", [])
        for yf in yml_files:
            path_from_root = yf.get("path_from_root", "")
            filepath = yf.get("filepath", "")
            _type = yf.get("label", "")
            role_info = yf.get("role_info", {})
            fr = FileResult(filepath=filepath, path_in_project=path_from_root, type=_type)
            fr.name_count = yf.get("name_count", 0)
            is_external_dependency = False
            if role_info:
                role_name = role_info["name"]
                role_path = role_info["path"]
                fr.role = role_name
                fr.role_path = role_path
                if role_info not in roles:
                    roles.append(role_info)
                is_external_dependency = role_info["is_external_dependency"]

            if _type == "error":
                error_info = yf.get("error", {})
                error_type = error_info.get("type", "")
                if error_type == "TooManyTasksError":
                    fr.skip_reason = SKIP_REASON.MANY_TASKS
                    fr.type = "others"
                elif error_type == "YAMLParseError":
                    fr.skip_reason = SKIP_REASON.YAMLParseError
                    fr.type = "others"
                else:
                    fr.error = error_type 
                out_scopes.append(fr)
            elif _type == "others":
                fr.skip_reason = SKIP_REASON.OTHERS
                out_scopes.append(fr)
            elif is_external_dependency:
                fr.skip_reason = SKIP_REASON.EXT_DEP
                out_scopes.append(fr)
            else:
                target_files.append(fr)

        # summarize loaded objects
        obj_types = {}
        for obj in objects:
            filepath = obj.get("filepath", "")
            _type = obj.get("type", "")
            name = obj.get("name", "")
            type_sum = obj_types.get(_type, [])
            type_sum.append({"filepath": filepath, "name": name})
            obj_types[_type] = type_sum

        obj_res = {}
        for t, objs in obj_types.items():
            obj_res[t] = {
                "count": len(objs),
                "objects": objs
            }

        file_results = []
        for result in target_files:
            tt = result.type
            fp = result.path_in_project
            result.in_scope = True            
            # file
            t_objs = obj_res.get(tt, {}).get("objects", [])
            is_scanned = False
            for to in t_objs:
                if to["filepath"] == fp:
                    is_scanned = True
            result.scanned = is_scanned
            if is_scanned:
                # task count
                if "task" in obj_res:
                    for tasks in obj_res["task"]["objects"]:
                        if fp == tasks["filepath"]:
                            result.scanned_task_count = result.scanned_task_count + 1
                if result.scanned_task_count == 0 and result.name_count != 0:
                    result.warning = "name found, but no task found"
            else:
                skip = False
                if result.scanned_task_count == 0 and result.name_count == 0:
                    result.in_scope = False
                    result.skip_reason = SKIP_REASON.NO_TASK
                    skip = True
                # todo: invalid taskfile/playbook
                elif result.scanned_task_count == 0 and result.role:
                    path_from_role = result.filepath.replace(result.role_path, "").lstrip("/")
                    if not path_from_role.startswith("tasks") and not path_from_role.startswith("handlers") \
                        and not path_from_role.startswith("tests"): 
                            result.in_scope = False
                            if result.type == "taskfile":
                                result.skip_reason = SKIP_REASON.INVALID_TASKFILE
                            elif result.type == "playbook":
                                result.skip_reason = SKIP_REASON.INVALID_PLAYBOOK
                            else:
                                result.skip_reason = f"invalid {result.type}"
                            skip = True
                if not skip:
                    result.error = "unknown"

            file_results.append(result)
        
        for result in out_scopes:
            result.in_scope = False
            result.scanned = False
            result.scanned_task_count = 0
            file_results.append(result)
        
        file_results = sorted(file_results, key=lambda x: x.path_in_project)

        # fail reason
        error_msgs = [res.error for res in file_results if res.error and not res.in_scope]
        p_error_msgs = [res.error for res in file_results if res.error and res.type == "playbook"]
        tf_error_msgs = [res.error for res in file_results if res.error and res.type == "taskfile"]
        error_msgs_count = dict(collections.Counter(error_msgs))
        p_error_msgs_count = dict(collections.Counter(p_error_msgs))
        tf_error_msgs_count = dict(collections.Counter(tf_error_msgs))
        p_skip_msgs = [res.skip_reason for res in file_results if res.skip_reason and res.type == "playbook"]
        tf_skip_msgs = [res.skip_reason for res in file_results if  res.skip_reason and res.type == "taskfile"]
        o_skip_msgs = [res.skip_reason for res in file_results if  res.skip_reason and res.type == "others"]
        p_skip_msgs_count = dict(collections.Counter(p_skip_msgs))
        tf_skip_msgs_count = dict(collections.Counter(tf_skip_msgs))
        o_skip_msgs_count = dict(collections.Counter(o_skip_msgs))
        warning_msgs = [res.warning for res in file_results if res.warning]
        warning_msgs_count = dict(collections.Counter(warning_msgs))
        
        # playbook
        p_total = sum(1 for res in file_results if res.type == "playbook")
        p_scanned = sum(1 for res in file_results if res.type == "playbook" and res.scanned)
        p_skipped = sum(1 for res in file_results if res.type == "playbook" and not res.in_scope)
        p_scan_error = sum(1 for res in file_results if res.type == "playbook" and res.in_scope and not res.scanned)
        p_skip_reasons = p_skip_msgs_count
        p_scan_error_msgs = p_error_msgs_count
        playbooks = ScanCount(p_total, p_scanned, p_skipped, p_skip_reasons, p_scan_error, p_scan_error_msgs)

        # taskfile
        t_total = sum(1 for res in file_results if res.type == "taskfile")
        t_scanned = sum(1 for res in file_results if res.type == "taskfile" and res.scanned)
        t_skipped = sum(1 for res in file_results if res.type == "taskfile" and not res.in_scope)
        t_scan_error = sum(1 for res in file_results if res.type == "taskfile" and res.in_scope and not res.scanned)
        t_skip_reasons = tf_skip_msgs_count
        t_scan_error_msgs = tf_error_msgs_count
        taskfiles = ScanCount(t_total, t_scanned, t_skipped, t_skip_reasons, t_scan_error, t_scan_error_msgs)

        # others
        o_total = sum(1 for res in file_results if res.type == "others")
        others = OtherCount(total=o_total, reason=o_skip_msgs_count)
    
        # errors
        errors = {}
        e_total = sum(1 for res in file_results if res.type == "error")
        error_reasons = error_msgs_count
        errors = ErrorCount(e_total, error_reasons)

        # total
        fc_total = len(file_results)
        scan_failures = t_scan_error + p_scan_error
        fc = FileCount(fc_total, scan_failures, playbooks, taskfiles, others, errors)

        # task count
        nc = sum(v.name_count for v in file_results)
        tt = sum(v.scanned_task_count for v in file_results)
        tc = TaskCount(total=tt, names=nc)

        # warnings
        warnings = warning_msgs_count

        # role
        rc = RoleCount(len(roles))

        # project src
        _source = objects[0].get("source", {}).get("type", "")
        _repo_name = objects[0].get("source", {}).get("repo_name", "")
        ps = ProjectSource(_source, _repo_name, object_json, meta_json)

        state = StateCount()
        if fc.playbooks.scan_error == 0 and fc.taskfiles.scan_error == 0 and fc.errors.total == 0:
            state.success = 1
        else:
            state.fail = 1

        # compute scan result
        scan_result = ScanReport(
            project_count=1,
            projects=[ps],
            file_count=fc,
            role_count=rc,
            task_count=tc,
            warning=warnings,
            state_count=state,
            file_results=file_results,
        )
        return scan_result
    
    def merge_repo_summary(self, repo_scan_results_path):
        # json_files = glob.glob(os.path.join(target_dir, "**", "summary.json"), recursive=True)
        results = load_json_data(repo_scan_results_path)
        merge_sr = ScanReport()
        for jd in results:
            sr = ScanReport.from_dict(jd)
            merge_sr = merge_sr.merge(sr)
        return merge_sr

    def merge_src_summary(self, target_dir):
        json_files = glob.glob(os.path.join(target_dir, "**", "summary.json"))
        merge_sr = ScanReport()
        for jf in json_files:
            if jf == os.path.join(target_dir, "summary.json"):
                continue
            jd = load_json_data(jf)
            sr = ScanReport.from_dict(jd[0])
            merge_sr = merge_sr.merge(sr)
        return merge_sr

    def merge_dicts(self, dict1, dict2):
        result_dict = {}
        for key in set(dict1.keys()).union(dict2.keys()):
            if key == "file_results":
                continue
            if isinstance(dict1.get(key), dict) and isinstance(dict2.get(key), dict):
                result_dict[key] = self.merge_dicts(dict1[key], dict2[key])
            else:
                # if type(dict1[key]) == int:
                value1 = dict1.get(key, 0)
                value2 = dict2.get(key, 0)
                # elif type(dict1[key]) == list:
                #     value1 = dict1.get(key, [])
                #     value2 = dict2.get(key, [])
                result_dict[key] = value1 + value2
        return result_dict
    
    def _gen_inventory_table(self, sr: ScanReport):
        # header = ["playbooks (scanned/total)", "taskfiles (scanned/total)", "others",  "ext. dependency", "parse error", "too many tasks", "total_num"]
        header = ["", "files", "scanned", "skipped", "error", "state", "skip_reasons", "error_reason"]
        c_num = len(header)
        r_num = 6
        cells = header
        # playbook
        state = ""
        if sr.file_count.playbooks.scan_error == 0:
            state = "✔"
        skip_reason = ""
        if sr.file_count.playbooks.skip_reasons:
            _r = []
            for key, count in sr.file_count.playbooks.skip_reasons.items():
                _r.append(f"{key} ({count})")
            skip_reason = ','.join(_r)
        err_reason = ""
        if sr.file_count.playbooks.scan_err_msgs:
            _r = []
            for key, count in sr.file_count.playbooks.scan_err_msgs.items():
                _r.append(f"{key} ({count})")
            err_reason = ','.join(_r)

        cells.append("playbook")
        cells.append(sr.file_count.playbooks.total)
        cells.append(sr.file_count.playbooks.scanned)
        cells.append(sr.file_count.playbooks.skipped)
        cells.append(sr.file_count.playbooks.scan_error)
        cells.append(state)
        cells.append(skip_reason)
        cells.append(err_reason)

        # taskfile
        state = ""
        if sr.file_count.taskfiles.scan_error == 0:
            state = "✔"
        skip_reason = ""
        if sr.file_count.taskfiles.skip_reasons:
            _r = []
            for key, count in sr.file_count.taskfiles.skip_reasons.items():
                _r.append(f"{key} ({count})")
            skip_reason = ','.join(_r)
        err_reason = ""
        if sr.file_count.taskfiles.scan_err_msgs:
            _r = []
            for key, count in sr.file_count.taskfiles.scan_err_msgs.items():
                _r.append(f"{key} ({count})")
            err_reason = ','.join(_r)

        cells.append("taskfile")
        cells.append(sr.file_count.taskfiles.total)
        cells.append(sr.file_count.taskfiles.scanned)
        cells.append(sr.file_count.taskfiles.skipped)
        cells.append(sr.file_count.taskfiles.scan_error)
        cells.append(state)
        cells.append(skip_reason)
        cells.append(err_reason)
        # others
        other_reason = ""
        if sr.file_count.others.reason:
            _r = []
            for key, count in sr.file_count.others.reason.items():
                _r.append(f"{key} ({count})")
            other_reason = ','.join(_r)
        cells.append("others")
        cells.append(sr.file_count.others.total)
        cells.append(0)
        cells.append(sr.file_count.others.total)
        cells.append(0)
        cells.append("✔")
        cells.append(other_reason)
        cells.append("")
        # errors
        err_reason = ""
        if sr.file_count.errors.err_msgs:
            _r = []
            for key, count in sr.file_count.errors.err_msgs.items():
                _r.append(f"{key} ({count})")
            err_reason = ','.join(_r)

        cells.append("errors")
        cells.append(sr.file_count.errors.total)
        cells.append(0)
        cells.append(0)
        cells.append(sr.file_count.errors.total)
        if sr.file_count.errors.total == 0:
            cells.append("✔")
        else:
            cells.append("")
        cells.append("")
        cells.append(err_reason)
        # total
        cells.append("Total")
        # files
        cells.append(sr.file_count.playbooks.total+sr.file_count.taskfiles.total+sr.file_count.others.total+sr.file_count.errors.total)
        # scanned
        cells.append(sr.file_count.playbooks.scanned+sr.file_count.taskfiles.scanned)
        # skipped
        cells.append(sr.file_count.playbooks.skipped+sr.file_count.taskfiles.skipped+sr.file_count.others.total)
        # error
        cells.append(sr.file_count.playbooks.scan_error+sr.file_count.taskfiles.scan_error+sr.file_count.errors.total)
        if sr.file_count.playbooks.scan_error+sr.file_count.taskfiles.scan_error+sr.file_count.errors.total == 0:
            cells.append("✔")
        else:
            cells.append("")
        cells.append("")
        cells.append("")
        return cells, c_num, r_num

    def _gen_project_result_table(self,  sr: ScanReport):
        header = ["passed", "error", "no result", "total"]
        c_num = len(header)
        r_num = 2
        cells = header
        cells.append(sr.state_count.success)
        cells.append(sr.state_count.fail)
        cells.append(sr.state_count.unknown)
        cells.append(sr.project_count)
        return cells, c_num, r_num

    def _gen_contents_summary_table(self, sr: ScanReport):
        header = ["playbooks", "taskfiles", "roles", "tasks"]
        c_num = len(header)
        r_num = 2
        cells = header
        cells.append(sr.file_count.playbooks.total)
        cells.append(sr.file_count.taskfiles.total)
        cells.append(sr.role_count.total)
        cells.append(sr.task_count.total)
        return cells, c_num, r_num

    def generate_repo_report(self, sr:ScanReport):
        # _data = load_json_data(json_path)
        # sr = ScanReport.from_dict(_data[0])
        state = self._get_state(sr)
        if state == "✔":
            return
        src_type = sr.projects[0].source
        repo_name = sr.projects[0].repo_name
        metadata_file = sr.projects[0].metadata_file
        object_file = sr.projects[0].object_file
        detail_repo_dir = os.path.join(self.md_outdir, src_type, repo_name)
        os.makedirs(detail_repo_dir, exist_ok=True)
        md_path = os.path.join(detail_repo_dir, "README.md")
        mdFile = MdUtils(file_name=md_path, title=f'Sage Scan Report : [{repo_name}](https://github.com/{repo_name})')

        top_md_path = os.path.join(self.outdir, "README.md")
        relative_path = os.path.relpath(top_md_path, detail_repo_dir)
        parent_md_path = os.path.join(self.md_outdir, src_type, "README.md")
        relative_parent_path = os.path.relpath(parent_md_path, detail_repo_dir)

        if sr.error:
            mdFile.new_line(f"[<< Top Report]({relative_path})")
            mdFile.new_line(f"[< Src Type Report]({relative_parent_path})")
            mdFile.new_header(level=1, title='Scan Error')
            mdFile.new_line(f"Error: {sr.error}")
            mdFile.new_line(f"metadata_file: {metadata_file}")
            mdFile.new_line(f"object_file: {object_file}")
            mdFile.create_md_file()
            return

        mdFile.new_line(f"[<< Top Report]({relative_path})")
        mdFile.new_line(f"[< Src Type Report]({relative_parent_path})")
        mdFile.new_line(f"metadata_file: {metadata_file}")
        mdFile.new_line(f"object_file: {object_file}")
        mdFile.new_header(level=1, title='Detail reports')
        mdFile.new_line("[Scan count](#scan-count)")
        mdFile.new_line("[Scan count (per Type)](#scan-count-per-type)")
        mdFile.new_line("[Scan count (per File)](#scan-count-per-file)")

        # mdFile.new_header(level=1, title='Scan count')
        # header = ["file count", "task count", "role count"]
        # cells = header
        # cells.append(sr.file_count.total)
        # cells.append(sr.task_count.total)
        # cells.append(sr.role_count.total)
        # mdFile.new_table(columns=3, rows=2, text=cells, text_align='left')
        mdFile.new_line('Scan count')
        cells, c_num, r_num = self._gen_contents_summary_table(sr)
        mdFile.new_table(columns=c_num, rows=r_num, text=cells, text_align='left')

        mdFile.new_header(level=1, title='Scan count (per Type)')
        cells, c_num, r_num = self._gen_inventory_table(sr)
        mdFile.new_table(columns=c_num, rows=r_num, text=cells, text_align='left')

        # link to description
        desc_md_path = os.path.join(self.outdir, "REASON-DESCRIPTION.md")
        relative_path = os.path.relpath(desc_md_path, detail_repo_dir)
        mdFile.new_paragraph(f'[skip reason]({relative_path})')


        mdFile.new_header(level=1, title='Scan count (per File)')
        header = ["file", "type", "status", "task count", "error", "skip_reason", "role", "warning"]
        cells = header
        for ty in sr.file_results:
            status = "" # success or fail or skipped or others
            if not ty.in_scope:
                status = "✔"
            elif ty.in_scope and not ty.scanned:
                status = ""
            elif ty.scanned:
                status = "✔"
            # cells.append(f"[{ty.path_in_project}](https://github.com/{repo_name}/blob/main/{ty.path_in_project})")
            cells.append(ty.path_in_project)
            cells.append(ty.type)
            cells.append(status)
            cells.append(ty.scanned_task_count)
            cells.append(ty.error)
            cells.append(ty.skip_reason)
            cells.append(ty.role)
            cells.append(ty.warning)
        mdFile.new_table(columns=8, rows=len(sr.file_results)+1, text=cells, text_align='left')
        mdFile.create_md_file()
        return
    
    def generate_src_report(self, src_type, src_summary_path, repo_scan_results_path):
        src_dir = os.path.join(self.md_outdir, src_type)
        md_path = os.path.join(src_dir, "README.md")
        os.makedirs(src_dir, exist_ok=True)
        mdFile = MdUtils(file_name=md_path, title=f'Sage Scan Report : {src_type}')

        top_md_path = os.path.join(self.outdir, "README.md")
        relative_path = os.path.relpath(top_md_path, src_dir)
        mdFile.new_line(f"[<< Top Report]({relative_path})")

        mdFile.new_line("")
        mdFile.new_line("[Scan Summary](#scan-summary)")
        mdFile.new_line("[Scan Summary (per Type)](#scan-summary-per-type)")
        mdFile.new_line("[Scan Summary (per Project)](#scan-summary-per-project)")

        # summary
        _data = load_json_data(src_summary_path)
        sr = ScanReport.from_dict(_data[0])

        # project
        mdFile.new_header(level=1, title='Scan Summary')
        mdFile.new_line('Project summary')
        cells, c_num, r_num = self._gen_project_result_table(sr)
        mdFile.new_table(columns=4, rows=2, text=cells, text_align='left')

        # role and task
        mdFile.new_line('Contents summary')
        cells, c_num, r_num = self._gen_contents_summary_table(sr)
        mdFile.new_table(columns=c_num, rows=r_num, text=cells, text_align='left')

        # ymls
        mdFile.new_header(level=1, title='Scan Summary (per Type)')
        cells, c_num, r_num = self._gen_inventory_table(sr)
        mdFile.new_table(columns=c_num, rows=r_num, text=cells, text_align='left')

        # link to description
        desc_md_path = os.path.join(self.outdir, "REASON-DESCRIPTION.md")
        relative_path = os.path.relpath(desc_md_path, src_dir)
        mdFile.new_paragraph(f'[skip reason]({relative_path})')

        # repo list
        mdFile.new_header(level=1, title='Scan Summary (per Project)')
        header = ["project", "state", "files", "playbooks (passed/total)", "taskfiles (passed/total)", "others", "skipped", "errors", "roles", "tasks", "warning"]
        cells = header
        # json_files = glob.glob(os.path.join(self.json_outdir, src_type, "**", "summary.json"), recursive=True)
        repo_results = load_json_data(repo_scan_results_path)

        passed_repos = []
        failed_repos = []
        for jd in repo_results:
            sr = ScanReport.from_dict(jd)
            state = self._get_state(sr)
            if state == "✔":
                passed_repos.append(sr)
            else:
                failed_repos.append(sr)

        row_count = 0
        for sr in failed_repos:
            row_count += 1
            src_type = sr.projects[0].source
            repo_name = sr.projects[0].repo_name
            repo_md_path = os.path.join(self.md_outdir, src_type, repo_name, "README.md")
            relative_path = os.path.relpath(repo_md_path, src_dir)

            state = self._get_state(sr)
            cells.append(f"[{repo_name}]({relative_path})")
            cells.append(state)
            cells.append(sr.file_count.total)
            cells.append(f"{sr.file_count.playbooks.scanned}/{sr.file_count.playbooks.total - sr.file_count.playbooks.skipped}")
            cells.append(f"{sr.file_count.taskfiles.scanned}/{sr.file_count.taskfiles.total - sr.file_count.taskfiles.skipped}")
            cells.append(sr.file_count.others.total)
            cells.append(sr.file_count.playbooks.skipped+sr.file_count.taskfiles.skipped)
            cells.append(sr.file_count.errors.total+sr.file_count.playbooks.scan_error+sr.file_count.taskfiles.scan_error)
            cells.append(sr.role_count.total)
            cells.append(sr.task_count.total)
            cells.append(sum(sr.warning.values()))

        for sr in passed_repos:
            row_count += 1
            src_type = sr.projects[0].source
            repo_name = sr.projects[0].repo_name
            repo_md_path = os.path.join(self.md_outdir, src_type, repo_name, "README.md")
            relative_path = os.path.relpath(repo_md_path, src_dir)

            state = self._get_state(sr)
            cells.append(repo_name)
            cells.append(state)
            cells.append(sr.file_count.total)
            cells.append(f"{sr.file_count.playbooks.scanned}/{sr.file_count.playbooks.total - sr.file_count.playbooks.skipped}")
            cells.append(f"{sr.file_count.taskfiles.scanned}/{sr.file_count.taskfiles.total - sr.file_count.taskfiles.skipped}")
            cells.append(sr.file_count.others.total)
            cells.append(sr.file_count.playbooks.skipped+sr.file_count.taskfiles.skipped)
            cells.append(sr.file_count.errors.total+sr.file_count.playbooks.scan_error+sr.file_count.taskfiles.scan_error)
            cells.append(sr.role_count.total)
            cells.append(sr.task_count.total)
            cells.append(sum(sr.warning.values()))

        mdFile.new_table(columns=11, rows=row_count+1, text=cells, text_align='left')
        mdFile.create_md_file()
        return

    def generate_top_report(self, json_path):
        md_path = os.path.join(self.outdir, "README.md")
        mdFile = MdUtils(file_name=md_path, title=f'Sage Scan Report')

        # summary
        _data = load_json_data(json_path)
        sr = ScanReport.from_dict(_data[0])

        mdFile.new_header(level=1, title='Scan Summary')
        mdFile.new_line('Project summary')
        cells, c_num, r_num = self._gen_project_result_table(sr)
        mdFile.new_table(columns=4, rows=2, text=cells, text_align='left')

        mdFile.new_line('Contents summary')
        cells, c_num, r_num = self._gen_contents_summary_table(sr)
        mdFile.new_table(columns=c_num, rows=r_num, text=cells, text_align='left')


        mdFile.new_header(level=1, title='Scan Summary (per Type)')
        cells, c_num, r_num = self._gen_inventory_table(sr)
        mdFile.new_table(columns=c_num, rows=r_num, text=cells, text_align='left')

        # link to description
        desc_md_path = os.path.join(self.outdir, "REASON-DESCRIPTION.md")
        relative_path = os.path.relpath(desc_md_path, self.outdir)
        mdFile.new_paragraph(f'[skip reason]({relative_path})')

        # src type result
        mdFile.new_header(level=1, title='Scan summary (per Src type)')
        header = ["src type", "state", "projects", "files", "playbooks (passed/total)", "taskfiles (passed/total)", "others", "skipped", "errors", "roles", "tasks"]
        cells = header

        json_files = glob.glob(os.path.join(self.json_outdir, "**", "summary.json"))
        row_count = 0
        for jf in json_files:
            row_count += 1
            _data = load_json_data(jf)
            sr = ScanReport.from_dict(_data[0])
            src_type = sr.projects[0].source
            repo_md_path = os.path.join(self.md_outdir, src_type, "README.md")
            relative_path = os.path.relpath(repo_md_path, self.outdir)
            
            state = self._get_state(sr)

            cells.append(f"[{src_type}]({relative_path})")
            cells.append(state)
            cells.append(sr.project_count)
            cells.append(sr.file_count.total)
            cells.append(f"{sr.file_count.playbooks.scanned}/{sr.file_count.playbooks.total - sr.file_count.playbooks.skipped}")
            cells.append(f"{sr.file_count.taskfiles.scanned}/{sr.file_count.taskfiles.total - sr.file_count.taskfiles.skipped}")
            cells.append(sr.file_count.others.total)
            cells.append(sr.file_count.playbooks.skipped+sr.file_count.taskfiles.skipped)
            cells.append(sr.file_count.errors.total+sr.file_count.playbooks.scan_error+sr.file_count.taskfiles.scan_error)
            cells.append(sr.role_count.total)
            cells.append(sr.task_count.total)
        mdFile.new_table(columns=11, rows=row_count+1, text=cells, text_align='left')
        mdFile.create_md_file()
        return

    def generate_reason_description(self):
        md_path = os.path.join(self.outdir, "REASON-DESCRIPTION.md")
        mdFile = MdUtils(file_name=md_path, title=f'Skip reason description')
        items = []
        items.append(f"{SKIP_REASON.NO_TASK}:  {SKIP_REASON.NO_TASK_desc}")
        items.append(f"{SKIP_REASON.INVALID_TASKFILE}:  {SKIP_REASON.INVALID_desc}")
        items.append(f"{SKIP_REASON.INVALID_PLAYBOOK}:  {SKIP_REASON.INVALID_desc}")
        items.append(f"{SKIP_REASON.EXT_DEP}:  {SKIP_REASON.EXT_DEP_desc}")
        items.append(f"{SKIP_REASON.OTHERS}:  {SKIP_REASON.OTHERS_desc}")
        items.append(f"{SKIP_REASON.MANY_TASKS}:  {SKIP_REASON.MANY_TASKS_desc}")
        items.append(f"{SKIP_REASON.YAMLParseError}:  {SKIP_REASON.YAMLParseError_desc}")
        mdFile.new_list(items)
        mdFile.create_md_file()
        return
    
    def _get_state(self, sr:ScanReport):
        state = ""
        if sr.state_count.fail != 0:
            state = ""
        elif sr.state_count.unknown != 0:
            state = "-"
        elif sr.state_count.fail == 0:
            state = "✔"
        return state

    def write_data(self, file):
        if len(self._report_buffer) > self.max_buffer:
            with open(file, mode="a") as f:
                for line in self._report_buffer:
                    f.write(f"{line.to_json(ensure_ascii=False).rstrip()}\n")
            self._report_buffer = []
        return
    
    def write_all(self, file):
        os.makedirs(os.path.join(self.json_outdir, src_type), exist_ok=True)
        file = os.path.join(self.json_outdir, src_type, "repo_scan_results.json")
        with open(file, mode="a") as f:
            for line in self._report_buffer:
                f.write(f"{line.to_json(ensure_ascii=False).rstrip()}\n")
        self._report_buffer = []
    

def export_result(filepath, results):
    with open(filepath, "w") as file:
        if type(results) == list:
            for result in results:
                json_str = jsonpickle.encode(result, make_refs=False, unpicklable=False)
                file.write(f"{json_str}\n")
        else:
            json_str = jsonpickle.encode(results, make_refs=False, unpicklable=False)
            file.write(f"{json_str}\n")

def load_json_data(filepath):
    with open(filepath, "r") as file:
        records = file.readlines()
    data = []
    for record in records:
        try:
            d = json.loads(record)
            data.append(d)
        except json.JSONDecodeError as e:
            print(f"JSONDecodeError: {e}")
    return data

def split_data(work_dir, obj_file, metadata_file):
    ds1 = Data_Splitter(work_dir, OBJ_FILE)
    with open(obj_file, "r") as f:
        for line in f:
            j_content = json.loads(line)
            source = j_content.get("source", {})
            src_type = source.get("type", None)
            repo_name = source.get("repo_name", None)
            ds1.save_to_buffer(src_type, repo_name, line)
    ds1.save_all()

    ds2 = Data_Splitter(work_dir, META_FILE)
    with open(metadata_file, "r") as f:
        for line in f:
            j_content = json.loads(line)
            source = j_content.get("source", {})
            src_type = source.get("type", None)
            repo_name = source.get("repo_name", None)
            ds2.save_to_buffer(src_type, repo_name, line)
    ds2.save_all()
    return


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="TODO")
    parser.add_argument("-o", "--out-dir", help='e.g. /tmp/batch/report')
    parser.add_argument("-i", "--input-dir", help="e.g. /tmp/batch/results")
    parser.add_argument("-f", "--file", help="object file")
    parser.add_argument("-m", "--metadata-file", help="metadata file")
    parser.add_argument("-t", "--type", help="e.g. GitHub-RHIBM")
    parser.add_argument("--from-gzip", action="store_true", help="read objects.json from gzip file")

    args = parser.parse_args()
    outdir = args.out_dir
    obj_file = args.file
    metadata_file = args.metadata_file
    input_dir = args.input_dir
    src_type = args.type
    gzip_input = args.from_gzip

    # if input are merged objects and metadata file, split data
    if obj_file and metadata_file and not input_dir:
        tmpdir = tempfile.TemporaryDirectory()
        tmp_dir = tmpdir.name
        split_data(tmp_dir, obj_file, metadata_file)
        objects_dir = tmp_dir
    
    if not os.path.exists(outdir):
        os.makedirs(outdir)
    json_outdir = os.path.join(outdir, "result_json")
    os.makedirs(json_outdir, exist_ok=True)
    report_dir = os.path.join(outdir, "detail")
    os.makedirs(report_dir, exist_ok=True)
    
    summarizer = ScanResultSummarizer(outdir, json_outdir, report_dir, input_dir)
    target_dir = os.path.join(summarizer.json_outdir, src_type)
    os.makedirs(target_dir, exist_ok=True)

    repo_scan_results_path = os.path.join(target_dir, "repo_scan_results.json")
    if gzip_input:
        files = get_target_files_from_gzip(input_dir, OBJ_FILE)
        with tarfile.open(input_dir, 'r:gz') as tar:
            summarizer.tar_file = tar
            summarizer.generate_repo_summary_all(files, repo_scan_results_path)
    else:
        objects_dir = os.path.join(input_dir, src_type)
        files = glob.glob(os.path.join(objects_dir, "**", OBJ_FILE), recursive=True)
        summarizer.generate_repo_summary_all(files, repo_scan_results_path)


    print("summarizing src_type scan result")
    src_type_summary = summarizer.merge_repo_summary(repo_scan_results_path)
    src_summary_path = os.path.join(target_dir, "summary.json")
    with open(src_summary_path, "w") as file:
        file.write(src_type_summary.to_json(ensure_ascii=False))
    # src_type md file
    summarizer.generate_src_report(src_type, src_summary_path, repo_scan_results_path)

    print("generating scan result for all")
    all_summary = summarizer.merge_src_summary(json_outdir)
    with open(os.path.join(json_outdir, "summary.json"), "w") as file:
        file.write(all_summary.to_json(ensure_ascii=False))
    # top md file
    summarizer.generate_top_report(os.path.join(json_outdir, "summary.json"))
    summarizer.generate_reason_description()
sage_scan/tools/object_report.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from mdutils.mdutils import MdUtils
from pathlib import Path
import os
import json
import argparse
import jsonpickle
from dataclasses import dataclass, field
import glob


@dataclass
class Summary():
    src_type: str = ''
    repo_name: str = ''
    findings: object = field(default_factory=dict)
    yml_inventory: object = field(default_factory=dict)
    sage_results: object = field(default_factory=dict)
    original: object = field(default_factory=dict)
    enriched: object = field(default_factory=dict)
    only_original: object = field(default_factory=dict)
    coverage: int = -1
    src_files: object = field(default_factory=dict)

def generate_summary(export_path, sage_dir, org_ft_dir, ftdata_src_type="", repo_name=""):
    scan_summary = Summary()

    scan_summary.src_type=ftdata_src_type
    scan_summary.repo_name=repo_name

    yml_inventory_file = os.path.join(sage_dir, "yml_inventory.json")
    findings_json = os.path.join(sage_dir, "findings.json")
    ftdata_file = os.path.join(org_ft_dir, "org-ftdata.json")
    sage_ftdata_file = os.path.join(sage_dir, "_tmp-ftdata.json")
    f_org_only = os.path.join(sage_dir, ftdata_src_type, "_only_org_ftdata.json")
    f_matched = os.path.join(sage_dir, ftdata_src_type, "_modified_ftdata.json")


    scan_summary.src_files = {
        "yml_inventory_file": yml_inventory_file,
        "findings_json": findings_json,
        "ftdata_file": ftdata_file,
        "sage_ftdata_file": sage_ftdata_file,
        "f_org_only": f_org_only,
        "f_matched": f_matched,
    }

    target_files, out_of_scopes = get_yml_inventory(yml_inventory_file)  # list: {path, label}
    
    tasks_in_findings = count_tasks_in_findings(findings_json)   # {path: task list}
    playbook_role = get_objects(findings_json)  # {playbook: [name], role: [name]}
    findings_task_count = sum(len(v) for v in tasks_in_findings.values())
    scan_summary.findings = {
        "task_count": findings_task_count,
        "playbooks": playbook_role["playbook"],
        "roles": playbook_role["role"],
        "file_tasks": tasks_in_findings,
    }

    playbook_count, taskfile_count, no_task_count, error_count, others_count, external_dep = get_file_type_count(out_of_scopes, target_files, tasks_in_findings)
    scan_summary.yml_inventory = {
        "playbook_count": playbook_count,
        "taskfile_count": taskfile_count,
        "no_task_count": no_task_count,
        "error_count": error_count,
        "others_count": others_count,
        "external_dep_count": external_dep,
        "target_files": target_files,
        "out_of_scopes": out_of_scopes,
    }

    task_count_org, o_playbooks, o_taskfiles = count_tasks_in_org_ftdata(ftdata_file)  #{path: task count}, p list, t list
    org_task_count = sum(task_count_org.values())
    scan_summary.original = {
        "task_count": org_task_count,
        "file_task_count": task_count_org,
        "playbooks": o_playbooks,
        "taskfiles": o_taskfiles,
        "playbook_count": len(o_playbooks),
        "taskfile_count": len(o_taskfiles),
    }

    task_count_sage, tasks_in_sage, s_playbooks, s_taskfiles = count_tasks_in_sage_ftdata(sage_ftdata_file) #{path: task count}, {path: task key list}, p list, t list
    sage_task_count = sum(task_count_sage.values())
    scan_summary.sage_results = {
        "task_count": sage_task_count,
        "file_task_count": task_count_sage,
        "file_task_keys": tasks_in_sage,
        "playbooks": s_playbooks,
        "taskfiles": s_taskfiles,
        "playbook_count": len(s_playbooks),
        "taskfile_count": len(s_taskfiles),
    }

    task_count_enriched, task_count_missing, missing_tasks = count_enriched_ftdata(f_org_only, f_matched) # {path: count}, list: {file, prompt}
    unchanged_sum = sum(task_count_missing.values())
    enriched_sum = sum(task_count_enriched.values())
    scan_summary.enriched = {
        "task_count": enriched_sum,
        "file_task_count": task_count_enriched
    }
    scan_summary.only_original = {
        "task_count": unchanged_sum,
        "file_task_count": task_count_missing,
        "only_org_detail": missing_tasks,
    }        

    if org_task_count != 0:
        p = enriched_sum/org_task_count*100
        coverage = round(p, 2)
        scan_summary.coverage = coverage

    export_result(export_path, scan_summary)
    return scan_summary

def get_yml_inventory(yml_inventory_file):
    target_files = []
    out_of_scope = []
    if os.path.exists(yml_inventory_file):
        with open(yml_inventory_file, 'r') as file:
            for line in file:
                item = json.loads(line)
                label = item["label"]
                role_info = item["role_info"]
                y = {
                    "path": item["path_from_root"].lstrip("/"),
                    "type": label,
                }
                if role_info:
                    role_name = role_info["name"]
                    y["role"] = role_name
                    is_external_dependency = role_info["is_external_dependency"]
                    if is_external_dependency:
                        y["external"] = True
                        out_of_scope.append(y)
                        continue
                if label == "others":
                    out_of_scope.append(y)
                    continue

                if label == "error":
                    out_of_scope.append(y)
                    continue
                target_files.append(y)
    # sort
    target_files = sorted(target_files, key=lambda x: x['type'])
    out_of_scope = sorted(out_of_scope, key=lambda x: x['type'])
    return target_files, out_of_scope

def count_tasks_in_findings(findings_json):
    file_tasks = {}
    if os.path.exists(findings_json):
        with open(findings_json, "r") as file:
            for line in file:
                findings = json.loads(line)
                mappings = findings.get("root_definitions", {}).get("mappings", {})
                target_name = mappings.get("target_name")

                definitions = findings.get("root_definitions", {}).get("definitions", {})
                playbooks = definitions.get("playbooks", [])
                for playbook in playbooks:
                    tasks = []
                    path = playbook["defined_in"].lstrip("/")
                    if not target_name.endswith(path):
                        path = f"{target_name}/{path}"
                    elif target_name != path:
                        path = target_name
                    plays = playbook.get("plays", [])
                    if len(plays) == 0:
                        if path not in file_tasks:
                            file_tasks[path] = []
                    else:
                        d_plays = definitions.get("plays", [])
                        for p in plays:
                            for dp in d_plays:
                                p_key = dp.get("key", "")
                                if p == p_key:
                                    p_tasks = dp.get("tasks", [])
                                    tasks.extend(p_tasks)
                                    post_tasks = dp.get("post_tasks", [])
                                    tasks.extend(post_tasks)
                                    pre_tasks = dp.get("pre_tasks", [])
                                    tasks.extend(pre_tasks)
                    uniq_tasks = list(set(tasks))
                    if path not in file_tasks:
                        file_tasks[path] = uniq_tasks
                    else:
                        _tasks = file_tasks[path]
                        _tasks.extend(uniq_tasks)
                        file_tasks[path] = list(set(_tasks))
                        
                taskfiles = definitions.get("taskfiles", [])
                for taskfile in taskfiles:
                    path = taskfile["defined_in"].lstrip("/")
                    if not target_name.endswith(path):
                        path = f"{target_name}/{path}"
                    elif target_name != path:
                        path = target_name
                    t_tasks = taskfile["tasks"]
                    if path not in file_tasks:
                        file_tasks[path] = t_tasks
                    else:
                        _tasks = file_tasks[path]
                        _tasks.extend(t_tasks)
                        file_tasks[path] = list(set(_tasks))
    # file_task_count = {}
    # for p, t in file_tasks.items():
    #     file_task_count[p] = len(t)
    return file_tasks

def get_objects(findings_json):
    objects = {"playbook": [], "role": []}
    if os.path.exists(findings_json):
        with open(findings_json, "r") as file:
            for line in file:
                findings = json.loads(line)
                definitions = findings.get("root_definitions", {}).get("definitions", {})
                playbooks = definitions.get("playbooks", [])
                # plays = definitions.get("plays", [])
                roles = definitions.get("roles", [])
                # taskfiles = definitions.get("taskfiles", [])
                for p in playbooks:
                    name = p.get("name", "")
                    if "playbook" in objects:
                        objects["playbook"].append(name)
                    else:
                        objects["playbook"] = [name]
                for r in roles:
                    name = r.get("name", "")
                    if "role" in objects:
                        objects["role"].append(name)
                    else:
                        objects["role"] = [name]
            for t, items in objects.items():
                objects[t] = list(set(items))
    return objects

def count_tasks_in_org_ftdata(ftdata_file):
    data = []
    if os.path.exists(ftdata_file):
        with open(ftdata_file, 'r') as file:
            for line in file:
                jo = json.loads(line)
                data.append(jo)
    task_count = {}
    playbooks = []
    taskfiles = []
    for d in data:
        path = d["path"].lstrip("/")
        task_count[path] = task_count.get(path, 0) + 1
        type = d["type"]
        if type == "playbook":
            if path not in playbooks:
                playbooks.append(path)
        if type == "task":
            if path not in taskfiles:
                taskfiles.append(path)
    return task_count, playbooks, taskfiles


def count_enriched_ftdata(f_org_only, f_matched):
    data = []
    if os.path.exists(f_org_only):
        with open(f_org_only, 'r') as file:
            for line in file:
                jo = json.loads(line)
                data.append(jo)
    missing_task_count = {}
    missing_tasks = []
    for d in data:
        path = d["path"].lstrip("/")
        missing_task_count[path] = missing_task_count.get(path, 0) + 1
        missing_tasks.append({
            "file": d["path"],
            "prompt": d["prompt"]
        })
    data = []
    if os.path.exists(f_matched):
        with open(f_matched, 'r') as file:
            for line in file:
                jo = json.loads(line)
                data.append(jo)
    matched_task_count = {}
    for d in data:
        path = d["path"].lstrip("/")
        matched_task_count[path] = matched_task_count.get(path, 0) + 1
    return matched_task_count, missing_task_count, missing_tasks

def count_tasks_in_sage_ftdata(sage_ftdata_file):
    data = []
    if os.path.exists(sage_ftdata_file):
        with open(sage_ftdata_file, 'r') as file:
            for line in file:
                jo = json.loads(line)
                data.append(jo)
    task_count = {}
    file_tasks = {}
    playbooks = []
    taskfiles = []
    for d in data:
        path = d["path"].lstrip("/")
        task_count[path] = task_count.get(path, 0) + 1
        if path not in file_tasks:
            file_tasks[path] = [d["ari_task_key"]]
        else:
            file_tasks[path].append(d["ari_task_key"])
        type = d["type"]
        if type == "playbook":
            if path not in playbooks:
                playbooks.append(path)
        if type == "taskfile":
            if path not in taskfiles:
                taskfiles.append(path)
    return task_count, file_tasks, playbooks, taskfiles

def get_file_type_count(out_of_scopes, target_files, tasks_in_findings):
    others_count = 0
    external_dep = 0
    error_count = 0
    no_task_count = 0
    playbook_count = 0
    taskfile_count = 0
    for oos in out_of_scopes:
        is_external = oos.get("external", False)
        type = oos["type"]
        if is_external:
            external_dep += 1
            continue
        if type == "others":
            others_count += 1
            continue
        if type == "error":
            error_count += 1
            continue
    for tf in target_files:
        type = tf["type"]
        path = tf["path"]
        fc = _get_task_count(path, tasks_in_findings)
        if fc == 0:
            no_task_count += 1
            continue
        if type == "playbook":
            playbook_count += 1
        if type == "taskfile":
            taskfile_count += 1
    return playbook_count, taskfile_count, no_task_count, error_count, others_count, external_dep

def _get_task_count(path, file_tasks):
    for p, t in file_tasks.items():
        if p.endswith(path):
            return len(t)
    return -1

def _get_tasks_in_findings(path, file_tasks):
    for p, t in file_tasks.items():
        if p.endswith(path):
            return t
    return []

def generate_detail_report(export_path, ss: Summary):
    if ss.repo_name:
        title = f"Sage Repo Scan Report: [{ss.repo_name}](https://github.com/{ss.repo_name})"
    else:
        title = f"Sage Repo Scan Report"
    mdFile = MdUtils(file_name=export_path, title=title)

    if src_type:
        mdFile.new_line(f"src_type: {src_type}")
        
    mdFile.new_line(f"yml_inventory_file: {ss.src_files['yml_inventory_file']}")
    mdFile.new_line(f"findings_json: {ss.src_files['findings_json']}")
    mdFile.new_line(f"original ftdata_file: {ss.src_files['ftdata_file']}")
    mdFile.new_line(f"sage_ftdata_file: {ss.src_files['sage_ftdata_file']}")

    mdFile.new_header(level=1, title='Detail reports')
    mdFile.new_line("[Yaml file inventory](#yaml-file-inventory)")
    mdFile.new_line("[Playbooks and Roles](#playbooks-and-roles)")
    mdFile.new_line("[Tasks scanned per file](#tasks-scanned-per-file)")
    mdFile.new_line("[Unrecorded tasks (parsed but rule failed)](#unrecorded-tasks)")
    mdFile.new_line("[Unchanged tasks (failed to add a new context)](#unchanged-tasks)")

    task_count_org = ss.original["task_count"]
    file_task_count_org = ss.original["file_task_count"]

    file_task_count_sage = ss.sage_results["file_task_count"]
    task_count_sage = ss.sage_results["task_count"]
    tasks_in_sage = ss.sage_results["file_task_keys"]
    s_playbook_count = ss.sage_results["playbook_count"]
    s_taskfile_count = ss.sage_results["taskfile_count"]

    task_count_missing = ss.only_original["task_count"]
    file_task_count_missing = ss.only_original["file_task_count"]
    missing_tasks = ss.only_original["only_org_detail"]
    file_task_count_enriched = ss.enriched["file_task_count"]
    
    findings_playbook = ss.findings["playbooks"]
    findings_role = ss.findings["roles"]
    findings_task_count = ss.findings["task_count"]
    tasks_in_findings = ss.findings["file_tasks"]

    playbook_count = ss.yml_inventory["playbook_count"]
    taskfile_count = ss.yml_inventory["taskfile_count"]
    no_task_count = ss.yml_inventory["no_task_count"]
    error_count = ss.yml_inventory["error_count"]
    others_count = ss.yml_inventory["others_count"]
    external_dep = ss.yml_inventory["external_dep_count"]
    target_files = ss.yml_inventory["target_files"]
    out_of_scopes = ss.yml_inventory["out_of_scopes"]


    mdFile.new_header(level=1, title='Yaml file inventory')
    header = ["playbooks", "taskfiles", "no task", "external dep", "others", "parse error", "total_num"]
    cells = header
    cells.append(playbook_count)
    cells.append(taskfile_count)
    cells.append(no_task_count)
    cells.append(external_dep)
    cells.append(others_count)
    cells.append(error_count)
    cells.append(playbook_count+taskfile_count+no_task_count+external_dep+others_count+error_count)
    mdFile.new_table(columns=7, rows=2, text=cells, text_align='left')
    
    # playbooks and roles
    mdFile.new_header(level=1, title='Playbooks and Roles')
    header = ["type", "name"]
    cells = header
    item_count = len(findings_playbook) + len(findings_role)
    for p in findings_playbook:
        cells.extend(["playbook", p])
    for r in findings_role:
        cells.extend(["role", r])
    mdFile.new_table(columns=2, rows=item_count+1, text=cells, text_align='left')

    # tasks
    header = ["path", "type", "status", "task count", "task recorded", "org tasks", "unchanged", "role"]
    cells = header
    status_count = {"fail": 0, "pass": 0}
    for ty in target_files:
        path = ty["path"]
        cells.append(path)
        cells.append(ty["type"])
        fc = _get_task_count(path, tasks_in_findings)
        if file_task_count_enriched.get(path, 0) < file_task_count_org.get(path, 0):
            cells.append("x")
            status_count["fail"] = status_count.get("fail", 0) + 1
        elif file_task_count_sage.get(path, 0) < fc:
            cells.append("△")
            status_count["fail"] = status_count.get("fail", 0) + 1
        else:
            cells.append("✔")
            status_count["pass"]= status_count.get("pass", 0) + 1
        cells.append(fc)
        cells.append(file_task_count_sage.get(path, 0))
        cells.append(file_task_count_org.get(path, 0))
        cells.append(file_task_count_missing.get(path, 0))
        cells.append(ty.get("role", ""))
    for o in out_of_scopes:
        cells.append(o["path"])
        cells.append(o["type"])
        cells.append("-")
        cells.append("-")
        cells.append("-")
        cells.append(file_task_count_org.get(o["path"], "-"))
        cells.append(file_task_count_missing.get(o["path"], "-"))
        cells.append(o.get("role", ""))
    # add summary table
    header = ["status count", "playbook (recorded/count)", "taskfile (recorded/count)", "task count", "task recorded", "org tasks", "unchanged"]
    cells_summary = header

    cells_summary.append(f"{status_count['pass']}/{status_count['pass']+status_count['fail']}")
    cells_summary.append(f"{s_playbook_count}/{playbook_count}")
    cells_summary.append(f"{s_taskfile_count}/{taskfile_count}")
    cells_summary.append(findings_task_count)
    cells_summary.append(task_count_sage)
    cells_summary.append(task_count_org)
    cells_summary.append(task_count_missing)
    mdFile.new_header(level=1, title='Summary')
    mdFile.new_table(columns=7, rows=2, text=cells_summary, text_align='left')

    # add tasks table
    mdFile.new_header(level=1, title='Tasks scanned per file')
    mdFile.new_table(columns=8, rows=len(target_files)+len(out_of_scopes)+1, text=cells, text_align='left')

    # unrecorded task (findings - sage ftdata)
    header = ["file", "task key"]
    mdFile.new_header(level=1, title='Unrecorded tasks')
    mdFile.new_line("List of tasks parsed but rule failed")
    cells = header
    count = 1
    for ty in target_files:
        path = ty["path"]
        f_task_keys = _get_tasks_in_findings(path, tasks_in_findings)
        s_task_keys = tasks_in_sage.get(path, [])
        for ft in f_task_keys:
            if ft not in s_task_keys:
                cells.append(path)
                cells.append(ft)
                count += 1
    mdFile.new_table(columns=2, rows=count, text=cells, text_align='left')

    # unchanged task (wisdom - sage ftdata)
    mdFile.new_header(level=1, title='Unchanged tasks')
    mdFile.new_line("List of tasks without new context")
    header = ["file", "prompt"]
    cells = header
    for t in missing_tasks:
         cells.append(t["file"])
         cells.append(t["prompt"])
    mdFile.new_table(columns=2, rows=len(missing_tasks)+1, text=cells, text_align='left')
    mdFile.create_md_file()
    return 

def generate_main_report(org_ftdata_dir, sage_dir, outdir, src_type, org_ftdata_type=""):
    if not org_ftdata_type:
        org_ftdata_type = src_type
    
    report_dir = os.path.join(outdir, org_ftdata_type)
    os.makedirs(report_dir, exist_ok=True)
    export_path = os.path.join(report_dir, "README.md")
    data_export_path = os.path.join(report_dir, "summary.json")
    mdFile = MdUtils(file_name=export_path, title='Sage Repo Scan Report')

    results = []
    target_sage_dir = os.path.join(sage_dir, src_type)
    for f_inventory in Path(target_sage_dir).rglob("yml_inventory.json"):
        if Path.is_file(f_inventory):
            f_inventory_str = str(f_inventory)
            parts = f_inventory_str.split("/")
            # project
            if parts[-3] != src_type:
                repo_name = f"{parts[-3]}/{parts[-2]}"
            else:
                repo_name = parts[-2]

            t_org_ftdata_dir = os.path.join(org_ftdata_dir, org_ftdata_type, repo_name)
            t_sage_dir = os.path.join(sage_dir, src_type, repo_name)

            detail_report_path = os.path.join(outdir, org_ftdata_type, repo_name, "README.md")
            detail_data_path = os.path.join(outdir, org_ftdata_type, repo_name, "summary.json")
            target_file_dir = os.path.dirname(detail_report_path)
            if not os.path.exists(target_file_dir):
                os.makedirs(target_file_dir)
            summary_data = generate_summary(detail_data_path, t_sage_dir, t_org_ftdata_dir, org_ftdata_type, repo_name)
            generate_detail_report(detail_report_path, summary_data)
            summary = {}
            summary["detail_report_path"] = detail_report_path
            summary["data"] = summary_data
            results.append(summary)

    sorted_results = sorted(results, key=lambda x: x["data"].coverage, reverse=True)

    mdFile.new_line(f"src_type: {sorted_results[0]['data'].src_type}")
    mdFile.new_line(f"src_file: RH_IBM_FT_data_GH_api.json")
    mdFile.new_line(f"original ftdata: {org_ftdata_dir}")

    json_data = []

    header = ["repo_name", "status", "playbooks (recorded/total)", "taskfiles", "roles", "tasks", "task recorded",  "enriched ftdata", "coverage (%)"]
    cells = header
    for res in sorted_results:
        ss = res["data"]
        src_type = ss.src_type
        repo_name = ss.repo_name
        detail_report_path = res["detail_report_path"]
        relative_path = os.path.relpath(detail_report_path, report_dir)
        cells.append(f"[{repo_name}]({relative_path})")
        p = ss.coverage
        status = "x"
        if p == 100 or p == -1:
            status = "✔"
        cells.append(status)

        playbook_count = ss.yml_inventory["playbook_count"]
        taskfile_count = ss.yml_inventory["taskfile_count"]
        playbook_count_sage = ss.sage_results["playbook_count"]
        taskfile_count_sage = ss.sage_results["taskfile_count"]

        task_count_org = ss.original["task_count"]
        task_count_sage = ss.sage_results["task_count"]
        task_count_findings = ss.findings["task_count"]
        task_count_en = ss.enriched["task_count"]

        cells.append(f'{playbook_count_sage}/{playbook_count}')
        cells.append(f'{taskfile_count_sage}/{taskfile_count}')
        cells.append(len(ss.findings["roles"]))
        cells.append(task_count_findings)
        cells.append(task_count_sage)
        cells.append(f'{task_count_en}/{task_count_org}')
        if p == -1:
            p = "N/A"
        cells.append(p)

        sj = {
            "src_type": src_type,
            "repo_name": repo_name,
            "status": status,
            "playbook_count_total": playbook_count,
            "playbook_count_sage": playbook_count_sage,
            "taskfile_count_total": taskfile_count,
            "taskfile_count_sage": taskfile_count_sage,
            "role_count": len(ss.findings["roles"]),
            "task_count_total": task_count_findings,
            "task_count_org": task_count_org,
            "task_count_sage": task_count_sage,
            "task_count_enriched": task_count_en,
            "coverage": ss.coverage,
            "detail_report_path": detail_report_path
        }
        json_data.append(sj)

    mdFile.new_table(columns=9, rows=len(sorted_results) + 1, text=cells, text_align='left')
    mdFile.create_md_file()
    json_result = {"md_path": export_path, "data": json_data}
    export_result(data_export_path, json_result)
    return

def export_result(filepath, results):
    with open(filepath, "w") as file:
        if type(results) == list:
            for result in results:
                json_str = jsonpickle.encode(result, make_refs=False, unpicklable=False)
                file.write(f"{json_str}\n")
        else:
            json_str = jsonpickle.encode(results, make_refs=False, unpicklable=False)
            file.write(f"{json_str}\n")

def load_json_data(filepath):
    with open(filepath, "r") as file:
        records = file.readlines()
    data = []
    for record in records:
        d = json.loads(record)
        data.append(d)
    return data

def generate_repo_src_type_report(st_report_dir, export_path, data_export_path, repo_results):
    sorted_repo_results = sorted(repo_results, key=lambda x: x["coverage"], reverse=True)
    mdFile = MdUtils(file_name=export_path, title=f'Sage Repo Scan Report - {sorted_repo_results[0]["repo_name"]}')    
    header = ["src_type", "status", "playbooks (recorded/total)", "taskfiles", "roles", "tasks", "task recorded",  "enriched ftdata", "coverage (%)"]
    cells = header
    for data in sorted_repo_results:
        relative_path = os.path.relpath(data["detail_report_path"], st_report_dir)
        cells.append(f"[{data['src_type']}]({relative_path})")
        cells.append(data["status"])
        cells.append(f"{data['playbook_count_sage']}/{data['playbook_count_total']}")
        cells.append(f"{data['taskfile_count_sage']}/{data['taskfile_count_total']}")
        cells.append(data["role_count"])
        cells.append(data["task_count_total"])
        cells.append(data["task_count_sage"])
        cells.append(f'{data["task_count_enriched"]}/{data["task_count_org"]}')
        if data["coverage"] == -1:
            cells.append("N/A")
        else:
            cells.append(data["coverage"])
    mdFile.new_table(columns=9, rows=len(sorted_repo_results) + 1, text=cells, text_align='left')
    mdFile.create_md_file()
    json_result = {"md_path": export_path, "data": sorted_repo_results}
    export_result(data_export_path, json_result)
    return

def generate_repo_summary_report(outdir, src_type, detail_outdir, subdir):
    files = glob.glob(os.path.join(subdir,f"{src_type}*","summary.json"))
    repo_data = {}
    for f in files:
        summary_data = load_json_data(f)
        data = summary_data[0]["data"]
        for result in data:
            repo_name = result["repo_name"]
            if repo_name in repo_data:
                repo_data[repo_name].append(result)
            else:
                repo_data[repo_name] = [result]

    repo_summary = []
    for repo, data in repo_data.items():
        coverage = -1
        status = "x"

        task_count_org = sum(v["task_count_org"] for v in data)
        task_count_enriched = sum(v["task_count_enriched"] for v in data)
        if task_count_org != 0:
            p = task_count_enriched/task_count_org*100
            coverage = round(p, 2)
        if coverage == 100 or coverage == -1:
            status = "✔"

        # generate detail md
        st_report_dir = os.path.join(detail_outdir, repo)
        os.makedirs(st_report_dir, exist_ok=True)
        detail_md_path = os.path.join(st_report_dir, "README.md")
        detail_data_path = os.path.join(st_report_dir, "summary.json")
        generate_repo_src_type_report(st_report_dir, detail_md_path, detail_data_path, data)

        result = {
            "repo_name": repo,
            "status": status,
            "playbook_count_total": data[0]["playbook_count_total"],
            "playbook_count_sage": data[0]["playbook_count_sage"],
            "taskfile_count_total": data[0]["taskfile_count_total"],
            "taskfile_count_sage": data[0]["taskfile_count_sage"],
            "role_count": data[0]["role_count"],
            "task_count_total": data[0]["task_count_total"],
            "task_count_org": task_count_org,
            "task_count_sage": data[0]["task_count_sage"],
            "task_count_enriched": task_count_enriched,
            "coverage": coverage,
            "detail_report_path": detail_md_path,
        }
        repo_summary.append(result)

    sorted_repo_summary = sorted(repo_summary, key=lambda x: x["coverage"], reverse=True)

    export_path = os.path.join(outdir, f"README-{src_type}-repo.md")
    data_export_path = os.path.join(outdir, f"summary-{src_type}-repo.json")
    mdFile = MdUtils(file_name=export_path, title='Sage Repo Scan Report')    
    header = ["repo_name", "status", "playbooks (recorded/total)", "taskfiles", "roles", "tasks", "task recorded",  "enriched ftdata", "coverage (%)"]
    cells = header
    for data in sorted_repo_summary:
        relative_path = os.path.relpath(data["detail_report_path"], outdir)
        cells.append(f'[{data["repo_name"]}]({relative_path})')
        cells.append(data["status"])
        cells.append(f"{data['playbook_count_sage']}/{data['playbook_count_total']}")
        cells.append(f"{data['taskfile_count_sage']}/{data['taskfile_count_total']}")
        cells.append(data["role_count"])
        cells.append(data["task_count_total"])
        cells.append(data["task_count_sage"])
        cells.append(f'{data["task_count_enriched"]}/{data["task_count_org"]}')
        if data["coverage"] == -1:
            cells.append("N/A")
        else:
            cells.append(data["coverage"])
    mdFile.new_table(columns=9, rows=len(sorted_repo_summary) + 1, text=cells, text_align='left')
    mdFile.create_md_file()
    json_result = {"md_path": export_path, "data": sorted_repo_summary}
    export_result(data_export_path, json_result)
    return

def generate_src_type_derivation_summary_report(_type, outdir, subdir):
    files = glob.glob(f'{subdir}/{_type}*/summary.json')
    repo_data = {}
    for f in files:
        summary_data = load_json_data(f)
        data = summary_data[0]["data"]
        for result in data:
            src_type = result["src_type"]
            if src_type in repo_data:
                repo_data[src_type].append(result)
            else:
                repo_data[src_type] = [result]

    repo_summary = []
    for src_type, data in repo_data.items():
        coverage = -1
        status = "x"

        task_count_org = sum(v["task_count_org"] for v in data)
        task_count_enriched = sum(v["task_count_enriched"] for v in data)
        if task_count_org != 0:
            p = task_count_enriched/task_count_org*100
            coverage = round(p, 2)
        if coverage == 100 or coverage == -1:
            status = "✔"

        result = {
            "src_type": src_type,
            "status": status,
            "playbook_count_total": sum(v["playbook_count_total"] for v in data),
            "playbook_count_sage": sum(v["playbook_count_sage"] for v in data),
            "taskfile_count_total": sum(v["taskfile_count_total"] for v in data),
            "taskfile_count_sage": sum(v["taskfile_count_sage"] for v in data),
            "role_count": sum(v["role_count"] for v in data),
            "task_count_total": sum(v["task_count_total"] for v in data),
            "task_count_org": task_count_org,
            "task_count_sage": sum(v["task_count_sage"] for v in data),
            "task_count_enriched": task_count_enriched,
            "coverage": coverage,
        }
        repo_summary.append(result)

    sorted_repo_summary = sorted(repo_summary, key=lambda x: x["coverage"], reverse=True)

    export_path = os.path.join(outdir, f"README-{_type}-st.md")
    data_export_path = os.path.join(outdir, f"summary-{_type}-st.json")
    mdFile = MdUtils(file_name=export_path, title=f'Sage Repo Scan Report {_type}')    
    header = ["src_type", "status", "playbooks (recorded/total)", "taskfiles", "roles", "tasks", "task recorded",  "enriched ftdata", "coverage (%)"]
    cells = header
    for data in sorted_repo_summary:
        cells.append(data["src_type"])
        cells.append(data["status"])
        cells.append(f"{data['playbook_count_sage']}/{data['playbook_count_total']}")
        cells.append(f"{data['taskfile_count_sage']}/{data['taskfile_count_total']}")
        cells.append(data["role_count"])
        cells.append(data["task_count_total"])
        cells.append(data["task_count_sage"])
        cells.append(f'{data["task_count_enriched"]}/{data["task_count_org"]}')
        if data["coverage"] == -1:
            cells.append("N/A")
        else:
            cells.append(data["coverage"])
    mdFile.new_table(columns=9, rows=len(sorted_repo_summary) + 1, text=cells, text_align='left')
    mdFile.create_md_file()
    json_result = {"md_path": export_path, "org_src_type": _type, "data": sorted_repo_summary}
    export_result(data_export_path, json_result)
    return

def generate_top_report(outdir):
    files = glob.glob(f'{outdir}/summary-*st.json')
    summary = []
    for file in files:
        jd = load_json_data(file)
        src_type = jd[0]["org_src_type"]
        data = jd[0]["data"]
        coverage = -1
        status = "x"

        task_count_org = sum(v["task_count_org"] for v in data)
        task_count_enriched = sum(v["task_count_enriched"] for v in data)
        if task_count_org != 0:
            p = task_count_enriched/task_count_org*100
            coverage = round(p, 2)
        if coverage == 100 or coverage == -1:
            status = "✔"

        result = {
            "src_type": src_type,
            "status": status,
            "playbook_count_total": data[0]["playbook_count_total"],
            "playbook_count_sage": data[0]["playbook_count_sage"],
            "taskfile_count_total": data[0]["taskfile_count_total"],
            "taskfile_count_sage": data[0]["taskfile_count_sage"],
            "role_count": data[0]["role_count"],
            "task_count_total": data[0]["task_count_total"],
            "task_count_org": task_count_org,
            "task_count_sage": data[0]["task_count_sage"],
            "task_count_enriched": task_count_enriched,
            "coverage": coverage,
        }
        summary.append(result)

    sorted_summary = sorted(summary, key=lambda x: x["coverage"], reverse=True)

    export_path = os.path.join(outdir, f"README.md")
    data_export_path = os.path.join(outdir, f"summary.json")
    mdFile = MdUtils(file_name=export_path, title=f'Sage Repo Scan Report')    
    header = ["src_type", "status", "playbooks (recorded/total)", "taskfiles", "roles", "tasks", "task recorded",  "enriched ftdata", "coverage (%)"]
    cells = header
    for data in sorted_summary:
        cells.append(data["src_type"])
        cells.append(data["status"])
        cells.append(f"{data['playbook_count_sage']}/{data['playbook_count_total']}")
        cells.append(f"{data['taskfile_count_sage']}/{data['taskfile_count_total']}")
        cells.append(data["role_count"])
        cells.append(data["task_count_total"])
        cells.append(data["task_count_sage"])
        cells.append(f'{data["task_count_enriched"]}/{data["task_count_org"]}')
        if data["coverage"] == -1:
            cells.append("N/A")
        else:
            cells.append(data["coverage"])
    mdFile.new_table(columns=9, rows=len(sorted_summary) + 1, text=cells, text_align='left')
    mdFile.create_md_file()
    json_result = {"md_path": export_path, "data": sorted_summary}
    export_result(data_export_path, json_result)
    return


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="TODO")
    # parser.add_argument("-f", "--file", help='md file')
    parser.add_argument("-o", "--out-dir", help='e.g. /tmp/batch/report')
    parser.add_argument("-s", "--sage-dir", help="e.g. /tmp/batch/results")
    parser.add_argument("--ft-data-dir", help="e.g. /tmp/batch/data")
    parser.add_argument("-t", "--type", help="e.g. GitHub-RHIBM")
    parser.add_argument("--ftdata-type", help='specify the type if original ftdata type is different from src-type (e.g. GitHub-RHIBM-disambiguate-module)')
    parser.add_argument("--single", action="store_true", help="")

    args = parser.parse_args()

    results = args.sage_dir
    ftdata_dir = args.ft_data_dir
    outdir = args.out_dir

    src_type = args.type
    org_ftdata_type = args.ftdata_type
    ftdata_src_type_list = [
        "",
        "-disambiguate-module", 
        "-disambiguate-module-prompt_mutation", 
        "-disambiguate-platform-or-module", 
        "-disambiguate-platform-or-module-prompt_mutation",
        "-prompt_mutation",
        "-prompt_mutation-prompt_mutation",
        ]

    if not os.path.isdir(outdir):
        os.makedirs(outdir)

    single_scan_mode = args.single

    if single_scan_mode:
        md_export_path = os.path.join(outdir, "report.md")
        json_export_path = os.path.join(outdir, "summary.json")
        summary_data = generate_summary(json_export_path, results, ftdata_dir)
        _ = generate_detail_report(md_export_path, summary_data)
    else:
        detail_dir = os.path.join(outdir, "details")
        if not os.path.exists(detail_dir):
            os.makedirs(detail_dir)
        # per repo in src-type
        for _derivation in ftdata_src_type_list:
            org_ftdata_type = f"{src_type}{_derivation}"
            generate_main_report(ftdata_dir, results, detail_dir, src_type, org_ftdata_type)

            # summarize report
            # per repo (repo_name.md)
            repo_dir = os.path.join(outdir, "repo-results")
            generate_repo_summary_report(outdir, src_type, repo_dir, detail_dir)
            # per src-type derivation (GitHub-RHIBM.md)
            generate_src_type_derivation_summary_report(src_type, outdir, detail_dir) 
            # all (README.md)
            generate_top_report(outdir)

sage_scan/tools/report.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
from dataclasses import dataclass, field
from dataclasses_json import dataclass_json
from typing import List


def dict_join(dict1, dict2):
    dic = dict1.copy()
    for k in dict2:
        if k in dic:
            dic[k] = dic[k] + dict2[k]
        else:
            dic[k] = dict2[k]
    return dic

@dataclass_json
@dataclass
class FileResult(object):
    filepath: str = ""
    path_in_project: str = ""
    type: str = ""
    role: str = ""
    role_path: str = ""
    error: str = ""
    skip_reason: str = ""
    in_scope: bool = False
    scanned: bool = False
    warning: str = ""
    name_count: int = 0
    scanned_task_count : int = 0
  

@dataclass_json
@dataclass
class ScanCount(object):
    total: int = 0
    scanned: int = 0
    skipped: int = 0
    skip_reasons: dict = field(default_factory=dict)
    scan_error: int = 0
    scan_err_msgs: dict = field(default_factory=dict)

    @classmethod
    def from_dict(cls, obj):
        c = cls()
        c.total = obj.get("total", 0)
        c.scanned = obj.get("scanned", 0)
        c.skipped = obj.get("skipped", 0)
        c.skip_reasons = obj.get("skip_reasons", {})
        c.scan_error = obj.get("scan_error", 0)
        c.scan_err_msgs = obj.get("scan_err_msgs", {})
        return c

    def merge(self, obj):
        if not isinstance(obj, ScanCount):
            raise Exception(f"incompatible type: {type(obj)}")
        obj2 = ScanCount()
        obj2.total = self.total + obj.total
        obj2.scanned = self.scanned + obj.scanned
        obj2.skipped = self.skipped + obj.skipped
        obj2.skip_reasons = dict_join(self.skip_reasons, obj.skip_reasons)
        obj2.scan_error = self.scan_error + obj.scan_error
        obj2.scan_err_msgs = dict_join(self.scan_err_msgs, obj.scan_err_msgs)
        return obj2


@dataclass_json
@dataclass
class OtherCount(object):
    total: int = 0
    reason: dict = field(default_factory=dict)

    @classmethod
    def from_dict(cls, obj):
        c = cls()
        c.total = obj.get("total", 0)
        c.reason = obj.get("reason", {})
        return c

    def merge(self, obj):
        if not isinstance(obj, OtherCount):
            raise Exception(f"incompatible type: {type(obj)}")
        obj2 = OtherCount()
        obj2.total = self.total + obj.total
        obj2.reason = dict_join(self.reason, obj.reason)
        return obj2


@dataclass_json
@dataclass
class ErrorCount(object):
    total: int = 0
    err_msgs: dict = field(default_factory=dict)

    @classmethod
    def from_dict(cls, obj):
        c = cls()
        c.total = obj.get("total", 0)
        c.err_msgs = obj.get("err_msgs", {})
        return c

    def merge(self, obj):
        if not isinstance(obj, ErrorCount):
            raise Exception(f"incompatible type: {type(obj)}")
        obj2 = ErrorCount()
        obj2.total = self.total + obj.total
        obj2.err_msgs = dict_join(self.err_msgs, obj.err_msgs)
        return obj2


@dataclass_json
@dataclass
class FileCount(object):
    total: int = 0
    scan_failures: int = 0
    playbooks: ScanCount = field(default_factory=ScanCount)
    taskfiles: ScanCount = field(default_factory=ScanCount)
    others: OtherCount = field(default_factory=OtherCount)
    errors: ErrorCount = field(default_factory=ErrorCount)

    @classmethod
    def from_dict(cls, obj):
        c = cls()
        c.total = obj.get("total", 0)
        c.scan_failures = obj.get("scan_failures", 0)
        c.playbooks = ScanCount.from_dict(obj.get("playbooks", {}))
        c.taskfiles = ScanCount.from_dict(obj.get("taskfiles", {}))
        c.others = OtherCount.from_dict(obj.get("others", {}))
        c.errors = ErrorCount.from_dict(obj.get("errors", {}))
        return c

    def merge(self, obj):
        if not isinstance(obj, FileCount):
            raise Exception(f"incompatible type: {type(obj)}")
        obj2 = FileCount()
        obj2.total = self.total + obj.total
        obj2.scan_failures = self.scan_failures + obj.scan_failures
        obj2.playbooks = self.playbooks.merge(obj.playbooks)
        obj2.taskfiles = self.taskfiles.merge(obj.taskfiles)
        obj2.others = self.others.merge(obj.others)
        obj2.errors = self.errors.merge(obj.errors)
        return obj2


@dataclass_json
@dataclass
class RoleCount(object):
    total: int = 0

    @classmethod
    def from_dict(cls, obj):
        c = cls()
        c.total = obj.get("total", 0)
        return c

    def merge(self, obj):
        if not isinstance(obj, RoleCount):
            raise Exception(f"incompatible type: {type(obj)}")
        obj2 = RoleCount()
        obj2.total = self.total + obj.total
        return obj2


@dataclass_json
@dataclass
class TaskCount(object):
    total: int = 0
    names: int = 0

    @classmethod
    def from_dict(cls, obj):
        c = cls()
        c.total = obj.get("total", 0)
        c.names = obj.get("names", 0)
        return c

    def merge(self, obj):
        if not isinstance(obj, TaskCount):
            raise Exception(f"incompatible type: {type(obj)}")
        obj2 = TaskCount()
        obj2.total = self.total + obj.total
        obj2.names = self.names + obj.names
        return obj2

@dataclass_json
@dataclass
class StateCount(object):
    success: int = 0
    fail: int = 0
    unknown: int = 0

    @classmethod
    def from_dict(cls, obj):
        c = cls()
        c.success = obj.get("success", 0)
        c.fail = obj.get("fail", 0)
        c.unknown = obj.get("unknown", 0)
        return c

    def merge(self, obj):
        if not isinstance(obj, StateCount):
            raise Exception(f"incompatible type: {type(obj)}")
        obj2 = StateCount()
        obj2.success = self.success + obj.success
        obj2.fail = self.fail + obj.fail
        obj2.unknown = self.unknown + obj.unknown
        return obj2

@dataclass_json
@dataclass
class ProjectSource(object):
    source: str = ""
    repo_name: str = ""
    object_file: str = ""
    metadata_file: str = ""

    @classmethod
    def from_dict(cls, obj):
        c = cls()
        c.metadata_file = obj.get("metadata_file", "")
        c.object_file = obj.get("object_file", "")
        c.repo_name = obj.get("repo_name", "")
        c.source = obj.get("source", "")
        return c

    @classmethod
    def from_list(cls, obj):
        ps_list = []
        if isinstance(obj, list):
            for v in obj:
                ps = cls.from_dict(v)
                ps_list.append(ps)
        return ps_list


@dataclass_json
@dataclass
class ScanReport(object):
    project_count: int = 0
    projects: List[ProjectSource] = field(default_factory=list)
    file_count: FileCount = field(default_factory=FileCount)
    role_count: RoleCount = field(default_factory=RoleCount)
    task_count: TaskCount = field(default_factory=TaskCount)
    warning: dict = field(default_factory=dict)
    error: dict = field(default_factory=dict)
    state_count: StateCount = field(default_factory=StateCount)
    file_results: List[FileResult] = field(default_factory=list)

    @classmethod
    def from_dict(cls, obj):
        c = cls()
        c.total = obj.get("project_count", 0)
        c.projects = ProjectSource.from_list(obj.get("projects", []))
        c.file_count = FileCount.from_dict(obj.get("file_count", {}))
        c.role_count = RoleCount.from_dict(obj.get("role_count", {}))
        c.task_count = TaskCount.from_dict(obj.get("task_count", {}))
        c.warning = obj.get("warning", {})
        c.error = obj.get("error", {})
        c.state_count = StateCount.from_dict(obj.get("state_count", {}))
        c.file_results = FileResult.from_dict(obj.get("file_results", []))
        return c

    def merge(self, obj):
        if not isinstance(obj, ScanReport):
            raise Exception(f"incompatible type: {type(obj)}")
        obj2 = ScanReport()
        obj2.project_count = self.project_count + obj.project_count
        obj2.projects = self.projects + obj.projects
        obj2.file_count = self.file_count.merge(obj.file_count)
        obj2.role_count = self.role_count.merge(obj.role_count)
        obj2.task_count = self.task_count.merge(obj.task_count)
        obj2.warning = dict_join(self.warning, obj.warning)
        obj2.error = dict_join(self.error, obj.error)
        obj2.state_count = self.state_count.merge(obj.state_count)
        return obj2


if __name__ == "__main__":
    file1 = "sage_scan/tools/summary-sample.json"
    file2 = "sage_scan/tools/summary-sample.json"

    d1 = {}
    d2 = {}

    with open(file1, "r") as f1:
        d1 = json.load(f1)

    with open(file1, "r") as f2:
        d2 = json.load(f2)

    sr1 = ScanReport.from_dict(d1)
    sr2 = ScanReport.from_dict(d2)

    sr_merged = sr1.merge(sr2)
    print(sr_merged.to_json(ensure_ascii=False))

sage_scan/tools/report_models.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
import jsonpickle
import argparse
from sage_scan.pipeline import SerializableRunContext
from ansible_risk_insight.models import (
    Module,
    Task,
    TaskCall,
    AnsibleRunContext,
)
from ansible_risk_insight.risk_assessment_model import RAMClient
import tabulate




ARI_KB_DATA_DIR = os.getenv("ARI_KB_DATA_DIR", None)

if ARI_KB_DATA_DIR is None:
    raise ValueError(f"Please specify an existing ARI KB dir by an env param:\n$ export ARI_KB_DATA_DIR=<PATH/TO/ARI_KB_DATA_DIR>")

if not os.path.exists(ARI_KB_DATA_DIR):
    raise ValueError(f"the ARI_KB_DATA_DIR does not exist: {ARI_KB_DATA_DIR}")

ram_client = RAMClient(root_dir=ARI_KB_DATA_DIR)


def get_module_fqcn(module_name):
    found_candidates = ram_client.search_module(module_name)
    if not found_candidates:
        return None

    # pick the best matched module instance here
    module = found_candidates[0]["object"]
    if not isinstance(module, Module):
        return None
    
    return module.fqcn


def main():
    parser = argparse.ArgumentParser(description="TODO")
    parser.add_argument("-f", "--file", help='result json file')
    args = parser.parse_args()

    fpath = args.file
    result = {}
    with open(fpath, "r") as file:
        body = file.read()
        result = jsonpickle.decode(body, safe=True)
    
    objects = result.get("objects", [])
    table_data = []
    not_found = []
    for obj in objects:
        if not isinstance(obj, SerializableRunContext):
            continue
        rc = obj.to_ansible_run_context()
        if not isinstance(rc, AnsibleRunContext):
            continue

        for node in rc:
            if not isinstance(node, TaskCall):
                continue
            if not isinstance(node.spec, Task):
                continue
            task = node.spec
            module_name = task.module
            module_fqcn = get_module_fqcn(module_name)
            if module_fqcn is None:
                not_found.append(module_name)
                continue
            table_data.append([module_name, module_fqcn])
        
    table_str = tabulate.tabulate(table_data, headers=["SHORT_NAME", "FQCN"])
    print(table_str)


if __name__ == '__main__':
    main()

sage_scan/tools/show_all_fqcn_from_result.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import json
import jsonpickle
import os


def prepare_source_dir(root_dir, yaml_file):
    path_list = []
    if not os.path.exists(root_dir):
        os.makedirs(root_dir)

    ignore_list = [
        "m4rkw/minotaur-install",
        "rodo/ansible-tsung",

        #GH2-PT
        "AutomateCompliance/AnsibleCompliancePlaybooks",
        "DevalexLLC/ansible-hardening-playbook",
        "ugns/ansible-ssg",
    ]
    
    yaml_file_contents = load_json_data(yaml_file)
    for content in yaml_file_contents:
        if "namespace_name" in content:
            type = "collection_role"
        else:
            type = "project"

        if type == "project":
            repo_name = content.get("repo_name")
            if repo_name in ignore_list:
                continue
            path = content.get("path")
            text = content.get("text")
            if not text:
                text = content.get("content")
            license = content.get("license")
            target_dir = os.path.join(root_dir, repo_name)
            if not os.path.exists(target_dir):
                os.makedirs(target_dir)
            target_file = os.path.join(target_dir, path.lstrip("/"))
            print(f"exporting yaml file {target_file}")
            target_file_dir = extract_directory(target_file)
            try:
                if not os.path.exists(target_file_dir):
                    os.makedirs(target_file_dir)
                with open(target_file, "w") as file:
                    file.write(text)
                path_list.append({
                    "repo_type": type,
                    "repo_name": repo_name,
                    "source": "",
                    "license": license,
                    "path": path,
                })
            except Exception as e:
                print(e)
        if type == "collection_role":
            namespace_name = content.get("namespace_name")
            path = content.get("path")
            if path == "":
                path = "example.yml"
            text = content.get("text")
            source = content.get("source")
            license = content.get("license")
            path_list.append({
                "repo_type": type,
                "repo_name": namespace_name,
                "source": source,
                "license": license,
                "path": path,
            })
            target_dir = os.path.join(root_dir, namespace_name)
            if not os.path.exists(target_dir):
                os.makedirs(target_dir)
            target_file = os.path.join(target_dir, path)
            target_file_dir = extract_directory(target_file)
            if not os.path.exists(target_file_dir):
                os.makedirs(target_file_dir)
            with open(target_file, "w") as file:
                # print(f"exporting yaml file {target_file}")
                file.write(text)
    return path_list

def extract_directory(file_path):
    directory = os.path.dirname(file_path)
    return directory

def load_json_data(filepath):
    with open(filepath, "r") as file:
        records = file.readlines()
    trains = []
    for record in records:
        train = json.loads(record)
        trains.append(train)
    return trains

def write_result(filepath, results):
    with open(filepath, "w") as file:
        if type(results) == list:
            for result in results:
                json_str = jsonpickle.encode(result, make_refs=False, unpicklable=False)
                file.write(f"{json_str}\n")
        else:
            json_str = jsonpickle.encode(results, make_refs=False, unpicklable=False)
            file.write(f"{json_str}\n")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="TODO")
    parser.add_argument("-t", "--source-type", help='source type (e.g."GitHub-RHIBM")')
    parser.add_argument("-s", "--source-json", help='source json file path (e.g. "/tmp/RH_IBM_FT_data_GH_api.json")')
    parser.add_argument("-o", "--out-dir", help="output directory")
    parser.add_argument("-p", "--project-list", help="project list")
    args = parser.parse_args()

    work_dir = args.out_dir
    src_type = args.source_type
    src_json = args.source_json
    project_list = args.project_list
    src_rb_dir = os.path.join(work_dir, "src_rb")
    path_list_dir = os.path.join(work_dir, "path_list")
    result_dir = os.path.join(work_dir, "results")

    os.makedirs(work_dir, exist_ok=True)
    os.makedirs(src_rb_dir, exist_ok=True)
    os.makedirs(path_list_dir, exist_ok=True)
    os.makedirs(result_dir, exist_ok=True)

    adir = os.path.join(src_rb_dir, src_type)
    if not os.path.exists(adir) or len(os.listdir(adir)) == 0:
        outfile = os.path.join(path_list_dir, f"path-list-{src_type}.txt")
        path_list = prepare_source_dir(adir, src_json)
        write_result(outfile, path_list)
sage_scan/tools/src_rebuilder.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path
import os
import json
import argparse
from treelib import Node, Tree
from ansible_risk_insight.findings import Findings
from ansible_risk_insight.models import Collection, Module, Playbook, Play, RoleInPlay, Repository, Role, TaskFile, Task
from ansible_risk_insight.keyutil import get_obj_type



types_map = {
    "collection": "collections",
    "module": "modules",
    "playbook": "playbooks",
    "play": "plays",
    "project": "projects",
    "role": "roles",
    "taskfile": "taskfiles",
    "task": "tasks",
}

# possible children types in spec objects
# NOTE: `playbook -> play -> task` is possible,
#       but `task -> (include) -> role` is not, because this checks only spec trees
tree_node_types = {
    "collections": ["playbooks", "roles", "modules", "plays", "taskfiles", "tasks"],
    "modules": [],
    "playbooks": ["plays", "tasks"],
    "plays": ["tasks"],
    "projects": ["collections", "playbooks", "roles", "modules", "plays", "taskfiles", "tasks"],
    "roles": ["playbooks", "plays", "taskfiles", "tasks", "modules"],
    "taskfiles": ["tasks"],
    "tasks": [],
}

def normalize_type_str(_type: str):
    
    lower_type = _type.lower()
    if lower_type in types_map:
        return types_map[lower_type]
    if lower_type in list(types_map.values()):
        return lower_type
    raise ValueError(f"Unknown ARI object type: {_type}")


def get_children_keys(obj):
    raw_children = obj.resolver_targets
    if not raw_children:
        return []
    children = []
    for c in raw_children:
        if isinstance(c, RoleInPlay):
            continue
        children.append(c)
    return children


def get_object_by_key(specs, key):
    type_str = get_obj_type(key)
    type_str = normalize_type_str(type_str)
    if type_str not in specs:
        return None
    objs = specs.get(type_str, [])
    found = None
    for obj in objs:
        if obj.key == key:
            found = obj
            break
    return found


def _recursive_get_children(specs, obj):
    obj_key = obj.key
    children_keys = get_children_keys(obj)
    if not children_keys:
        return []
    
    children = []
    for c_key in children_keys:
        c_obj = get_object_by_key(specs, c_key)
        c_name = ""
        if c_obj:
            c_name = create_node_name_from_obj(c_obj)
        else:
            c_name = "(Unknown)"
        node = (c_name, c_key, obj_key)
        children.append(node)

        if c_obj:
            offspring = _recursive_get_children(specs, c_obj)
            children.extend(offspring) 

    return children


def create_node_name_from_obj(obj):
    _type = ""
    _name = ""
    if isinstance(obj, Collection):
        _type = "Collection"
        _name = obj.fqcn
    elif isinstance(obj, Module):
        _type = "Module"
        _name = obj.fqcn
    elif isinstance(obj, Playbook):
        _type = "Playbook"
        _name = obj.defined_in
    elif isinstance(obj, Play):
        _type = "Play"
        _name = obj.name
    elif isinstance(obj, Repository):
        _type = "Project"
        _name = obj.name
    elif isinstance(obj, Role):
        _type = "Role"
        _name = obj.fqcn
        if obj.metadata and isinstance(obj.metadata, dict):
            role_desc = obj.metadata.get("galaxy_info", {}).get("description", "")
            if role_desc:
                _name = _name + f" ({role_desc})"
    elif isinstance(obj, TaskFile):
        _type = "TaskFile"
        _name = obj.defined_in
    elif isinstance(obj, Task):
        _type = "Task"
        _name = obj.name
    if not _name:
        _name = "(No name)"
    name = f"{_type} {_name}"
    return name


def create_object_tree_by_specs(specs, _type):
    if not specs:
        return None
    
    if not isinstance(specs, dict):
        return None
    
    # create root nodes
    trees = []
    for obj in specs.get(_type, []):
        key = obj.key
        tree = Tree()
        root_name = create_node_name_from_obj(obj)
        tree.create_node(root_name, key)
        tree_nodes = _recursive_get_children(specs, obj)

        for node in tree_nodes:
            node_name, obj_key, parent_key = node
            tree.create_node(node_name, obj_key, parent=parent_key)

        trees.append(tree)
    
    for tree in trees:
        tree.show(sorting=False)
        print("-" * 90)



if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="TODO")
    parser.add_argument("-f", "--file", help='findings json file')
    parser.add_argument("-t", "--type", help='type of objects to view')
    # parser.add_argument("-d", "--dir", help='tmp dir to recreate source dir')
    # parser.add_argument("-o", "--out-file", help="output directory for the rule evaluation result")
    args = parser.parse_args()

    fpath = args.file
    _type = normalize_type_str(args.type)
    findings_list = []
    with open(fpath, "r") as file:
        for line in file:
            f = Findings.load(json_str=line)
            findings_list.append(f)
    
    for findings in findings_list:
        specs = findings.root_definitions.get("definitions", {})
        if _type not in specs:
            continue
        
        create_object_tree_by_specs(specs, _type)
        




sage_scan/tools/view_findings.py
IBM/sage
# -*- mode:python; coding:utf-8 -*-

# Copyright (c) 2023 IBM Corp. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from ansible_risk_insight.scanner import ARIScanner, config, Config
from ansible_risk_insight.models import NodeResult, RuleResult
from ansible_risk_insight.finder import (
    find_all_ymls,
    label_yml_file,
    get_role_info_from_path,
    get_project_info_for_file,
)
from ansible_risk_insight.utils import escape_local_path
from ansible_risk_insight.risk_detector import load_rules
import os
import pathlib
import pygit2
import argparse
import time
import traceback
import joblib
import threading
import jsonpickle
import json


def get_git_version():
    version = ""
    try:
        repo_dir = pathlib.Path(__file__).parent.parent.resolve()
        repo = pygit2.Repository(repo_dir)
        version = str(repo.head.target)
    except Exception:
        pass
    return version


def get_yml_label(file_path, root_path):
    relative_path = file_path.replace(root_path, "")
    if relative_path[-1] == "/":
        relative_path = relative_path[:-1]
    
    label, error = label_yml_file(file_path)
    role_name, role_path = get_role_info_from_path(file_path)
    role_info = None
    if role_name and role_path:
        role_info = {"name": role_name, "path": role_path}

    project_name, project_path = get_project_info_for_file(file_path, root_path)
    project_info = None
    if project_name and project_path:
        project_info = {"name": project_name, "path": project_path}
    
    # print(f"[{label}] {relative_path} {role_info}")
    if error:
        print(f"failed to get yml label:\n {error}")
        label = "error"
    return label, role_info, project_info


def create_scan_list(yml_inventory):
    role_file_list = {}
    project_file_list = {}
    independent_file_list = []
    for yml_data in yml_inventory:
        filepath = yml_data["filepath"]
        path_from_root = yml_data["path_from_root"]
        label = yml_data["label"]
        role_info = yml_data["role_info"]
        in_role = yml_data["in_role"]
        project_info = yml_data["project_info"]
        in_project = yml_data["in_project"]
        if project_info:
            p_name = project_info.get("name", "")
            p_path = project_info.get("path", "")
            if p_name not in project_file_list:
                project_file_list[p_name] = {"path": p_path, "files": []}
            project_file_list[p_name]["files"].append({
                "filepath": filepath,
                "path_from_root": path_from_root,
                "label": label,
                "project_info": project_info,
                "role_info": role_info,
                "in_project": in_project,
                "in_role": in_role,
            })
        elif role_info:
            r_name = role_info.get("name", "")
            r_path = role_info.get("path", "")
            if role_info.get("is_external_dependency", False):
                continue
            if r_name not in role_file_list:
                role_file_list[r_name] = {"path": r_path, "files": []}
            role_file_list[r_name]["files"].append({
                "filepath": filepath,
                "path_from_root": path_from_root,
                "label": label,
                "project_info": project_info,
                "role_info": role_info,
                "in_project": in_project,
                "in_role": in_role,
            })
        else:
            independent_file_list.append({
                "filepath": filepath,
                "path_from_root": path_from_root,
                "label": label,
                "project_info": project_info,
                "role_info": role_info,
                "in_project": in_project,
                "in_role": in_role,
            })
    return project_file_list, role_file_list, independent_file_list



def get_yml_list(root_dir: str):
    found_ymls = find_all_ymls(root_dir)
    all_files = []
    for yml_path in found_ymls:
        label, role_info, project_info = get_yml_label(yml_path, root_dir)
        if not role_info:
            role_info = {}
        if not project_info:
            project_info = {}
        if role_info:
            if role_info["path"] and not role_info["path"].startswith(root_dir):
                role_info["path"] = os.path.join(root_dir, role_info["path"])
            role_info["is_external_dependency"] = True if "." in role_info["name"] else False
        in_role = True if role_info else False
        in_project = True if project_info else False
        all_files.append({
            "filepath": yml_path,
            "path_from_root": yml_path.replace(root_dir, "").lstrip("/"),
            "label": label,
            "role_info": role_info,
            "project_info": project_info,
            "in_role": in_role,
            "in_project": in_project,
        })
    return all_files


def get_rule_id_list(rules_dir=""):
    rules = load_rules(rules_dir=rules_dir)
    rule_id_list = [rule.rule_id for rule in rules]
    return rule_id_list


def load_new_fidnings_list(fpath):
    new_findings_list = []
    with open(fpath, "r") as file:
        for line in file:
            new_findings = jsonpickle.decode(line)
            new_findings_list.append(new_findings)
            break
    return new_findings_list


def load_new_fidnings(fpath):
    _list = load_new_fidnings_list(fpath)
    if not _list:
        return None
    return _list[0]


def dump_new_findings_list(new_findings_list, fpath):
    lines = []
    for new_findings in new_findings_list:
        new_findings_json_str = jsonpickle.encode(new_findings, make_refs=False)
        lines.append(new_findings_json_str + "\n")
        break

    out_dir = os.path.dirname(fpath)
    if not os.path.exists(out_dir):
        os.makedirs(out_dir, exist_ok=True)

    with open(fpath, "w") as outfile:
        outfile.write("".join(lines))
sage_scan/utils.py
IBM/sage