custom.py

# --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------

import os
import sys
import tempfile
from typing import Optional, BinaryIO

from azext_confcom import oras_proxy, os_util, security_policy
from azext_confcom._validators import resolve_stdio
from azext_confcom.config import (
    DEFAULT_REGO_FRAGMENTS, POLICY_FIELD_CONTAINERS_ELEMENTS_REGO_FRAGMENTS,
    REGO_IMPORT_FILE_STRUCTURE, ACI_FIELD_VERSION, ACI_FIELD_CONTAINERS)
from azext_confcom.cose_proxy import CoseSignToolProxy
from azext_confcom.errors import eprint
from azext_confcom.fragment_util import get_all_fragment_contents
from azext_confcom.init_checks import run_initial_docker_checks
from azext_confcom.kata_proxy import KataPolicyGenProxy
from azext_confcom.security_policy import AciPolicy, OutputType
from azext_confcom.template_util import (
    get_image_name, inject_policy_into_template, inject_policy_into_yaml,
    pretty_print_func, print_existing_policy_from_arm_template,
    print_existing_policy_from_yaml, print_func, str_to_sha256)
from azext_confcom.command.fragment_attach import fragment_attach as _fragment_attach
from azext_confcom.command.fragment_push import fragment_push as _fragment_push
from azext_confcom.command.containers_from_image import containers_from_image as _containers_from_image
from azext_confcom.command.containers_from_vn2 import containers_from_vn2 as _containers_from_vn2
from knack.log import get_logger
from pkg_resources import parse_version

logger = get_logger(__name__)


# pylint: disable=too-many-locals, too-many-branches
def acipolicygen_confcom(
    input_path: str,
    arm_template: str,
    arm_template_parameters: str,
    image_name: str,
    virtual_node_yaml_path: str,
    infrastructure_svn: str,
    tar_mapping_location: str,
    container_definitions: Optional[list] = None,
    approve_wildcards: str = False,
    outraw: bool = False,
    outraw_pretty_print: bool = False,
    diff: bool = False,
    validate_sidecar: bool = False,
    save_to_file: str = None,
    debug_mode: bool = False,
    print_policy_to_terminal: bool = False,
    disable_stdio: Optional[bool] = None,
    enable_stdio: Optional[bool] = None,
    print_existing_policy: bool = False,
    faster_hashing: bool = False,
    omit_id: bool = False,
    include_fragments: bool = False,
    fragments_json: str = None,
    exclude_default_fragments: bool = False,
):
    if print_existing_policy or outraw or outraw_pretty_print:
        logger.warning(
            "%s %s %s %s %s",
            "Secrets that are included in the provided arm template or configuration files",
            "in the container env or cmd sections will be printed out with this flag.",
            "These are outputed secrets that you must protect. Be sure that you do not include these secrets in your",
            "source control. Also verify that no secrets are present in the logs of your command or script.",
            "For additional information, see http://aka.ms/clisecrets. \n",
        )

    if container_definitions is None:
        container_definitions = []

    stdio_enabled = resolve_stdio(enable_stdio, disable_stdio)

    if print_existing_policy and arm_template:
        print_existing_policy_from_arm_template(arm_template, arm_template_parameters)
        return
    if print_existing_policy and virtual_node_yaml_path:
        print_existing_policy_from_yaml(virtual_node_yaml_path)
        return

    if debug_mode:
        logger.warning("WARNING: %s %s",
                       "Debug mode must only be used for debugging purposes. ",
                       "It should not be used for production systems.\n")

    tar_mapping = tar_mapping_validation(tar_mapping_location)

    output_type = get_output_type(outraw, outraw_pretty_print)

    container_group_policies = None

    # warn user that input infrastructure_svn is less than the configured default value
    check_infrastructure_svn(infrastructure_svn)

    fragments_list = []
    # gather information about the fragments being used in the new policy
    if include_fragments:
        fragments_data = os_util.load_json_from_file(fragments_json or input_path)
        if isinstance(fragments_data, dict):
            fragments_list = fragments_data.get("fragments", [])
            # standalone fragments from external file
            fragments_list.extend(fragments_data.get("standaloneFragments", []))

        # convert to list if it's just a dict. if it's empty, make it an empty list
        if not fragments_data:
            fragments_list = []
        elif not isinstance(fragments_list, list):
            fragments_list = [fragments_list]

    # telling the user what operation we're doing
    logger.warning(
        "Generating security policy for %s: %s in %s",
        "ARM Template" if arm_template else "Image" if image_name else "Input File",
        input_path or arm_template or image_name or virtual_node_yaml_path,
        "base64"
        if output_type == security_policy.OutputType.DEFAULT
        else "clear text",
    )
    # error checking for making sure an input is provided is above
    if input_path:
        container_group_policies = security_policy.load_policy_from_json_file(
            input_path,
            debug_mode=debug_mode,
            infrastructure_svn=infrastructure_svn,
            disable_stdio=(not stdio_enabled),
            exclude_default_fragments=exclude_default_fragments,
        )
    elif arm_template:
        container_group_policies = security_policy.load_policy_from_arm_template_file(
            infrastructure_svn,
            arm_template,
            arm_template_parameters,
            debug_mode=debug_mode,
            disable_stdio=(not stdio_enabled),
            approve_wildcards=approve_wildcards,
            diff_mode=diff,
            rego_imports=fragments_list,
            exclude_default_fragments=exclude_default_fragments,
        )
    elif image_name:
        container_group_policies = security_policy.load_policy_from_image_name(
            image_name, debug_mode=debug_mode, disable_stdio=(not stdio_enabled)
        )
    elif virtual_node_yaml_path:
        container_group_policies = security_policy.load_policy_from_virtual_node_yaml_file(
            virtual_node_yaml_path=virtual_node_yaml_path,
            debug_mode=debug_mode,
            disable_stdio=(not stdio_enabled),
            approve_wildcards=approve_wildcards,
            diff_mode=diff,
            rego_imports=fragments_list,
            exclude_default_fragments=exclude_default_fragments,
            infrastructure_svn=infrastructure_svn,
        )
    elif container_definitions:
        container_group_policies = AciPolicy(
            {
                ACI_FIELD_VERSION: "1.0",
                ACI_FIELD_CONTAINERS: [],
            },
            debug_mode=debug_mode,
            disable_stdio=disable_stdio,
            container_definitions=container_definitions,
        )

    exit_code = 0

    # standardize the output so we're only operating on arrays
    # this makes more sense than making the "from_file" and "from_image" outputting arrays
    # since they can only ever output a single image's policy
    if not isinstance(container_group_policies, list):
        container_group_policies = [container_group_policies]

    # get all of the fragments that are being used in the policy
    # and associate them with each container group

    if include_fragments:
        logger.info("Including fragments in the policy")
        fragment_policy_list = []
        container_names = []
        fragment_imports = []
        for policy in container_group_policies:
            fragment_imports.extend(policy.get_fragments())
            for container in policy.get_images():
                container_names.append(container.get_container_image())
        # get all the fragments that are being used in the policy
        fragment_policy_list = get_all_fragment_contents(container_names, fragment_imports)
        for policy in container_group_policies:
            policy.set_fragment_contents(fragment_policy_list)

    for count, policy in enumerate(container_group_policies):
        # this is where parameters and variables are populated
        policy.populate_policy_content_for_all_images(
            individual_image=bool(image_name), tar_mapping=tar_mapping, faster_hashing=faster_hashing
        )

        if validate_sidecar:
            exit_code = validate_sidecar_in_policy(policy, output_type == security_policy.OutputType.PRETTY_PRINT)
        elif virtual_node_yaml_path and not (print_policy_to_terminal or outraw or outraw_pretty_print or diff):
            result = inject_policy_into_yaml(
                virtual_node_yaml_path, policy.get_serialized_output(omit_id=omit_id), count
            )
            if result:
                print(str_to_sha256(policy.get_serialized_output(OutputType.RAW, omit_id=omit_id)))
                logger.info("CCE Policy successfully injected into YAML file")
        elif diff:
            exit_code = get_diff_outputs(policy, output_type == security_policy.OutputType.PRETTY_PRINT)
        elif arm_template and not (print_policy_to_terminal or outraw or outraw_pretty_print):
            result = inject_policy_into_template(arm_template, arm_template_parameters,
                                                 policy.get_serialized_output(omit_id=omit_id), count)
            if result:
                # this is always going to be the unencoded policy
                print(str_to_sha256(policy.get_serialized_output(OutputType.RAW, omit_id=omit_id)))
                logger.info("CCE Policy successfully injected into ARM Template")

        else:
            # output to terminal
            print(f"{policy.get_serialized_output(output_type, omit_id=omit_id)}\n\n")

            # output to file
            if save_to_file:
                logger.warning(
                    "%s %s %s",
                    "(Deprecation Warning) the --save-to-file (-s) flag is deprecated ",
                    "and will be removed in a future release. ",
                    "Please print to the console and redirect to a file instead."
                )
                policy.save_to_file(save_to_file, output_type)

    if exit_code != 0:
        sys.exit(exit_code)


# pylint: disable=R0914
def acifragmentgen_confcom(
    image_name: str,
    input_path: str,
    tar_mapping_location: str,
    namespace: str,
    svn: str,
    feed: str,
    key: str,
    chain: str,
    minimum_svn: str,
    container_definitions: Optional[list] = None,
    image_target: str = "",
    algo: str = "ES384",
    fragment_path: str = None,
    omit_id: bool = False,
    generate_import: bool = False,
    disable_stdio: Optional[bool] = None,
    enable_stdio: Optional[bool] = None,
    debug_mode: bool = False,
    output_filename: str = "",
    outraw: bool = False,
    upload_fragment: bool = False,
    no_print: bool = False,
    fragments_json: str = "",
    out_signed_fragment: bool = False,
):
    if container_definitions is None:
        container_definitions = []

    stdio_enabled = resolve_stdio(enable_stdio, disable_stdio)

    output_type = get_fragment_output_type(outraw)

    if generate_import:
        cose_client = CoseSignToolProxy()
        import_statements = []
        # images can have multiple fragments attached to them so we need an array to hold the import statements
        if fragment_path:
            # download and cleanup the fragment from registry if it's not local already
            downloaded_fragment = False
            if not os.path.exists(fragment_path):
                fragment_path = oras_proxy.pull(fragment_path)
                downloaded_fragment = True
            import_statements = [cose_client.generate_import_from_path(fragment_path, minimum_svn=minimum_svn)]
            if downloaded_fragment:
                os_util.clean_up_temp_folder(fragment_path)

        elif image_name:
            import_statements = oras_proxy.generate_imports_from_image_name(image_name, minimum_svn=minimum_svn)

        fragments_file_contents = {}
        fragments_list = []
        if fragments_json:
            logger.info("Creating/appending import statement JSON file")
            if os.path.isfile(fragments_json):
                fragments_file_contents = os_util.load_json_from_file(fragments_json)
                if isinstance(fragments_file_contents, list):
                    eprint(
                        "%s %s %s %s",
                        "Unsupported JSON file format. ",
                        "Please make sure the outermost structure is not an array. ",
                        "An empty import file should look like: ",
                        REGO_IMPORT_FILE_STRUCTURE,
                        exit_code=1
                    )

                fragments_list = fragments_file_contents.get(POLICY_FIELD_CONTAINERS_ELEMENTS_REGO_FRAGMENTS, [])

        # convert to list if it's just a dict
        if isinstance(fragments_list, dict):
            fragments_list = [fragments_list]
        fragments_list += import_statements

        fragments_file_contents[POLICY_FIELD_CONTAINERS_ELEMENTS_REGO_FRAGMENTS] = fragments_list
        if fragments_json:
            os_util.write_str_to_file(fragments_json, pretty_print_func(fragments_file_contents))
        else:
            print(pretty_print_func(fragments_file_contents))
        return

    tar_mapping = tar_mapping_validation(tar_mapping_location, using_config_file=bool(input_path))

    if image_name:
        policy = security_policy.load_policy_from_image_name(
            image_name, debug_mode=debug_mode, disable_stdio=(not stdio_enabled)
        )
    elif input_path:
        # this is using --input
        if not tar_mapping:
            tar_mapping = os_util.load_tar_mapping_from_config_file(input_path)
        policy = security_policy.load_policy_from_json_file(
            input_path, debug_mode=debug_mode, disable_stdio=(not stdio_enabled)
        )
    elif container_definitions:
        policy = AciPolicy(
            {
                ACI_FIELD_VERSION: "1.0",
                ACI_FIELD_CONTAINERS: [],
            },
            debug_mode=debug_mode,
            disable_stdio=disable_stdio,
            container_definitions=container_definitions,
        )
    else:
        eprint("Either --image-name, --input, or --container-definitions must be provided", exit_code=2)
        return

    # get all of the fragments that are being used in the policy
    # and associate them with each container group
    fragment_policy_list = []
    container_names = []
    fragment_imports = policy.get_fragments()
    for container in policy.get_images():
        container_names.append(container.get_container_image())
    fragment_policy_list = get_all_fragment_contents(container_names, fragment_imports)
    policy.set_fragment_contents(fragment_policy_list)
    policy.populate_policy_content_for_all_images(
        individual_image=bool(image_name), tar_mapping=tar_mapping
    )

    # make sure we have images to generate a fragment
    policy_images = policy.get_images()
    if not policy_images and not container_definitions:
        eprint("No images found in the policy or all images are covered by fragments")

    if not feed:
        # strip the tag or hash off the image name so there are stable feed names
        feed = get_image_name(image_target)

    fragment_text = policy.generate_fragment(namespace, svn, output_type, omit_id=omit_id)

    if output_type != security_policy.OutputType.DEFAULT and not no_print and not out_signed_fragment:
        print(fragment_text)

    # take ".rego" off the end of the filename if it's there, it'll get added back later
    output_filename = output_filename.replace(".rego", "")
    filename = f"{output_filename or namespace}.rego"

    if out_signed_fragment:
        filename = os.path.join(tempfile.gettempdir(), filename)

    os_util.write_str_to_file(filename, fragment_text)

    if key:
        cose_proxy = CoseSignToolProxy()
        iss = cose_proxy.create_issuer(chain)
        out_path = filename + ".cose"

        if out_signed_fragment:
            out_path = os.path.join(tempfile.gettempdir(), os.path.basename(out_path))

        cose_proxy.cose_sign(filename, key, chain, feed, iss, algo, out_path)

        # Preserve default behaviour established since version 1.1.0 of attaching
        # the fragment to the first image specified in input
        # (or --image-target if specified)
        if upload_fragment:
            oras_proxy.attach_fragment_to_image(
                image_name=image_target or policy_images[0].containerImage,
                filename=out_path,
            )

        if out_signed_fragment:
            with open(out_path, "rb") as f:
                sys.stdout.buffer.write(f.read())


def katapolicygen_confcom(
    yaml_path: str,
    config_map_file: str,
    outraw: bool = False,
    print_policy: bool = False,
    use_cached_files: bool = False,
    settings_file_name: str = None,
    rules_file_name: str = None,
    print_version: bool = False,
    containerd_pull: str = False,
    containerd_socket_path: str = None,
):
    kata_proxy = KataPolicyGenProxy()

    output = kata_proxy.kata_genpolicy(
        yaml_path,
        config_map_file=config_map_file,
        outraw=outraw,
        print_policy=print_policy,
        use_cached_files=use_cached_files,
        settings_file_name=settings_file_name,
        rules_file_name=rules_file_name,
        print_version=print_version,
        containerd_pull=containerd_pull,
        containerd_socket_path=containerd_socket_path,
    )
    print(output)


def update_confcom(cmd, instance, tags=None):
    with cmd.update_context(instance) as c:
        c.set_param("tags", tags)
    return instance


def check_infrastructure_svn(infrastructure_svn):
    if infrastructure_svn and parse_version(infrastructure_svn) < parse_version(
        DEFAULT_REGO_FRAGMENTS[0]["minimum_svn"]
    ):
        logger.warning(
            "Input Infrastructure Fragment Software Version Number is less than the default Infrastructure SVN: %s",
            DEFAULT_REGO_FRAGMENTS[0]["minimum_svn"],
        )


def validate_sidecar_in_policy(policy: security_policy.AciPolicy, outraw_pretty_print: bool):
    is_valid, output = policy.validate_sidecars()

    if outraw_pretty_print:
        formatted_output = pretty_print_func(output)
    else:
        formatted_output = print_func(output)

    if is_valid:
        print("Sidecar containers will pass with its generated policy")
        return 0

    print(
        f"Sidecar containers will not pass with its generated policy: {formatted_output}"
    )
    return 2


def get_diff_outputs(policy: security_policy.AciPolicy, outraw_pretty_print: bool):
    exit_code = 0
    is_valid, output = policy.validate_cce_policy()

    if outraw_pretty_print:
        formatted_output = pretty_print_func(output)
    else:
        formatted_output = print_func(output)

    print(
        "Existing policy and Template match"
        if is_valid
        else formatted_output
    )
    fragment_diff = policy.compare_fragments()

    if fragment_diff != {}:
        logger.warning(
            "Fragments in the existing policy are not the defaults. If this is expected, ignore this warning."
        )
    if not is_valid:
        logger.warning(
            "Existing Policy and Template differ. Consider recreating the base64-encoded policy."
        )
        exit_code = 2
    return exit_code


# TODO: refactor this function to use _validators.py functions and make sure the tar path
# isn't coming from the config file rather than the flag
def tar_mapping_validation(tar_mapping_location: str, using_config_file: bool = False):
    tar_mapping = None
    if tar_mapping_location:
        if not os.path.isfile(tar_mapping_location):
            print(
                "--tar input must either be a path to a json file with "
                + "image to tar location mappings or the location to a single tar file."
            )
            sys.exit(2)
        # file is mapping images to tar file locations
        elif tar_mapping_location.endswith(".json"):
            tar_mapping = os_util.load_tar_mapping_from_file(tar_mapping_location)
        # passing in a single tar location for a single image policy
        else:
            tar_mapping = tar_mapping_location
    elif not using_config_file:
        # only need to do the docker checks if we're not grabbing image info from tar files
        error_msg = run_initial_docker_checks()
        if error_msg:
            logger.warning(error_msg)
            sys.exit(1)
    return tar_mapping


def get_output_type(outraw, outraw_pretty_print):
    output_type = security_policy.OutputType.DEFAULT
    if outraw:
        output_type = security_policy.OutputType.RAW
    elif outraw_pretty_print:
        output_type = security_policy.OutputType.PRETTY_PRINT
    return output_type


def get_fragment_output_type(outraw):
    output_type = security_policy.OutputType.PRETTY_PRINT
    if outraw:
        output_type = security_policy.OutputType.RAW
    return output_type


def fragment_attach(
    signed_fragment: BinaryIO,
    manifest_tag: str,
) -> None:
    _fragment_attach(
        signed_fragment=signed_fragment,
        manifest_tag=manifest_tag
    )


def fragment_push(
    signed_fragment: BinaryIO,
    manifest_tag: str,
) -> None:
    _fragment_push(
        signed_fragment=signed_fragment,
        manifest_tag=manifest_tag
    )


def containers_from_image(
    image: str,
    platform: str,
) -> None:
    _containers_from_image(
        image=image,
        platform=platform,
    )


def containers_from_vn2(
    template: str,
    container_name: Optional[str] = None,
) -> None:
    print(_containers_from_vn2(
        template=template,
        container_name=container_name,
    ))