def _install()

in setup/setup-ci/security-scanner/amlsecscan.py [0:0]
59 lines of code
8 McCabe index (conditional complexity)

def _install(log_analytics_resource_id):
    if os.geteuid() != 0:
        raise Exception(
            "Installation must be performed by the root user. Please run again using sudo."
        )

    _logger.debug(f"Creating folder {_config_folder_path}")
    os.makedirs(_config_folder_path, exist_ok=True)
    shutil.chown(_config_folder_path, "azureuser", "azureuser")

    config = {"logAnalyticsResourceId": None}

    # Load config file if present
    if os.path.exists(_local_config_path):
        _logger.debug(f"Loading configuration from {_local_config_path}")
        with open(_local_config_path, "rt") as file:
            config.update(json.load(file))
        _logger.debug(
            f"logAnalyticsResourceId after loading config file: {config['logAnalyticsResourceId']}"
        )

    # Set Log Analytics workspace ARM Resource ID if passed via command-line parameter
    if log_analytics_resource_id is not None:
        config["logAnalyticsResourceId"] = log_analytics_resource_id
        _logger.debug(
            f"logAnalyticsResourceId after setting command-line parameter: {config['logAnalyticsResourceId']}"
        )

    # Retrieve Log Analytics workspace ARM Resource ID from Azure ML diagnostic settings if
    # provided neither via local config file nor command-line parameter
    if config.get("logAnalyticsResourceId", None) is None:
        config["logAnalyticsResourceId"] = _get_log_analytics_from_diagnostic_settings()
        _logger.debug(
            f"logAnalyticsResourceId after querying Azure ML diagnostic settings: {config['logAnalyticsResourceId']}"
        )

    # Sanitize the Log Analytics workspace ARM Resource ID
    config["logAnalyticsResourceId"] = _sanitize_log_analytics_resource_id(
        config["logAnalyticsResourceId"]
    )

    _logger.debug(f"Configuration: {config}")

    _logger.info(f"Writing configuration file {_global_config_path}")
    with open(_global_config_path, "wt") as file:
        json.dump(config, file, indent=2)
    shutil.chown(_global_config_path, "azureuser", "azureuser")

    _logger.info("Installing Trivy")
    _run(
        "apt-get install -y --no-install-recommends --quiet wget apt-transport-https gnupg lsb-release"
    )
    _run(
        "wget -qO - https://aquasecurity.github.io/trivy-repo/deb/public.key | apt-key add -"
    )
    _run(
        "echo deb https://aquasecurity.github.io/trivy-repo/deb $(lsb_release -sc) main | tee -a /etc/apt/sources.list.d/trivy.list"
    )
    _run("apt-get update")
    _run("apt-get install -y --no-install-recommends --quiet trivy")

    script_path = _config_folder_path + "/run.sh"
    _logger.info(f"Writing script file {script_path}")
    with open(script_path, "wt") as file:
        file.write(
            f"""#!/bin/bash
set -e
exec 1> >(logger -s -t AMLSECSCAN) 2>&1

# Limit CPU usage to 20% and reduce priority (note: the configuration is not persisted during reboot)
if [ ! -d /sys/fs/cgroup/cpu/amlsecscan ]
then
    mkdir -p /sys/fs/cgroup/cpu/amlsecscan
    echo 100000 | tee /sys/fs/cgroup/cpu/amlsecscan/cpu.cfs_period_us > /dev/null
    echo 20000 | tee /sys/fs/cgroup/cpu/amlsecscan/cpu.cfs_quota_us > /dev/null
    echo 5 | tee /sys/fs/cgroup/cpu/amlsecscan/cpu.shares > /dev/null
fi
echo $$ | tee /sys/fs/cgroup/cpu/amlsecscan/tasks > /dev/null

nice -n 19 python3 {os.path.abspath(__file__)} $1 $2 $3 $4 $5
"""
        )
    os.chmod(script_path, 0o0755)

    _logger.info(f"Writing crontab file /etc/cron.d/amlsecscan")
    with open("/etc/cron.d/amlsecscan", "wt") as file:
        file.write(
            f"""*/10 * * * * root {script_path} heartbeat
37 5 * * * root {script_path} scan all
@reboot root sleep 600 && {script_path} scan all
"""
        )
    os.chmod("/etc/cron.d/amlsecscan", 0o0644)