#!/usr/bin/env pypy3
# Microsoft Azure Linux Agent
#
# Copyright 2018 Microsoft Corporation
#
# 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
#
#     http://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.
#

# This script forces the process check by putting unknown process in the agent's cgroup

import subprocess
import datetime

from assertpy import fail

from azurelinuxagent.common.future import UTC
from azurelinuxagent.common.utils import shellutil
from tests_e2e.tests.lib.cgroup_helpers import check_agent_quota_disabled, check_log_message, get_unit_cgroup_proc_path, AGENT_SERVICE_NAME
from tests_e2e.tests.lib.logging import log
from tests_e2e.tests.lib.retry import retry_if_false


def prepare_agent():
    check_time = datetime.datetime.now(UTC)
    log.info("Executing script update-waagent-conf to enable agent cgroups config flag")
    result = shellutil.run_command(["update-waagent-conf", "Debug.CgroupCheckPeriod=20", "Debug.CgroupLogMetrics=y",
                                    "Debug.CgroupDisableOnProcessCheckFailure=y",
                                    "Debug.CgroupDisableOnQuotaCheckFailure=n"])
    log.info("Successfully enabled agent cgroups config flag: {0}".format(result))

    found: bool = retry_if_false(lambda: check_log_message(" Agent cgroups enabled: True", after_timestamp=check_time))
    if not found:
        fail("Agent cgroups not enabled")


def creating_dummy_process():
    log.info("Creating dummy process to add to agent's cgroup")
    dd_command = ["sleep", "60m"]
    proc = subprocess.Popen(dd_command)
    return proc.pid


def remove_dummy_process(pid):
    log.info("Removing dummy process from agent's cgroup")
    shellutil.run_command(["kill", "-9", str(pid)])


def disable_agent_cgroups_with_unknown_process(pid):
    """
    Adding dummy process to the agent's cgroup and verifying that the agent detects the unknown process and disables cgroups

    Note: System may kick the added process out of the cgroups, keeps adding until agent detect that process
    """

    def unknown_process_found():
        cgroup_procs_path = get_unit_cgroup_proc_path(AGENT_SERVICE_NAME, 'cpu,cpuacct')
        log.info("Adding dummy process %s to cgroup.procs file %s", pid, cgroup_procs_path)
        try:
            with open(cgroup_procs_path, 'a') as f:
                f.write("\n")
                f.write(str(pid))
        except Exception as e:
            log.warning("Error while adding process to cgroup.procs file: {0}".format(e))
            return False

        # The log message indicating the check failed is similar to
        #     2021-03-29T23:33:15.603530Z INFO MonitorHandler ExtHandler Disabling resource usage monitoring. Reason: Check on cgroups failed:
        #     [CGroupsException] The agent's cgroup includes unexpected processes: ['[PID: 25826] python3\x00/home/nam/Compute-Runtime-Tux-Pipeline/dungeon_crawler/s']
        found: bool = retry_if_false(lambda: check_log_message(
            "Disabling resource usage monitoring. Reason: Check on cgroups failed:.+The agent's cgroup includes unexpected processes:.+{0}".format(
                pid)), attempts=3)
        return found and retry_if_false(check_agent_quota_disabled, attempts=3)

    found: bool = retry_if_false(unknown_process_found, attempts=3)
    if not found:
        fail("The agent did not detect unknown process: {0}".format(pid))


def main():
    prepare_agent()
    pid = creating_dummy_process()
    disable_agent_cgroups_with_unknown_process(pid)
    remove_dummy_process(pid)


if __name__ == "__main__":
    main()
