#!/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 checks firewall rules are set on boot through cron job logs.And also capture the logs for debugging purposes.
#
import argparse
import os
import re
import shutil

from assertpy import fail

from azurelinuxagent.common.utils import shellutil
from tests_e2e.tests.lib.firewall_helpers import verify_all_rules_exist
from tests_e2e.tests.lib.logging import log
from tests_e2e.tests.lib.retry import retry


def move_cron_logs_to_var_log():
    # Move the cron logs to /var/log
    log.info("Moving cron logs to /var/log for debugging purposes")
    for cron_log in [ROOT_CRON_LOG, NON_ROOT_CRON_LOG, NON_ROOT_WIRE_XML, ROOT_WIRE_XML]:
        try:
            shutil.move(src=cron_log, dst=os.path.join("/var", "log",
                                                       "{0}.{1}".format(os.path.basename(cron_log),
                                                                        BOOT_NAME)))
        except Exception as e:
            log.info("Unable to move cron log to /var/log; {0}".format(e))


def check_wireserver_versions_file_exist(wire_version_file):
    log.info("Checking wire-versions file exist: {0}".format(wire_version_file))
    if not os.path.exists(wire_version_file):
        log.info("File: {0} not found".format(wire_version_file))
        return False

    if os.stat(wire_version_file).st_size > 0:
        return True

    return False


def verify_data_in_cron_logs(cron_log, verify, err_msg):
    log.info("Verifying Cron logs")

    def cron_log_checks():

        if not os.path.exists(cron_log):
            raise Exception("Cron log file not found: {0}".format(cron_log))
        with open(cron_log) as f:
            cron_logs_lines = list(map(lambda _: _.strip(), f.readlines()))
        if not cron_logs_lines:
            raise Exception("Empty cron file, looks like cronjob didnt run")

        if any("Unable to connect to network, exiting now" in line for line in cron_logs_lines):
            raise Exception("VM was unable to connect to network on startup. Skipping test validation")

        if not any("ExitCode" in line for line in cron_logs_lines):
            raise Exception("Cron logs still incomplete, will try again in a minute")

        if not any(verify(line) for line in cron_logs_lines):
            fail("Verification failed! (UNEXPECTED): {0}".format(err_msg))

        log.info("Verification succeeded. Cron logs as expected")

    retry(cron_log_checks)


def verify_wireserver_ip_reachable_for_root():
    """
    For root logs -
    Ensure the /var/log/wire-versions-root.xml is not-empty (generated by the cron job)
    Ensure the exit code in the /var/log/reboot-cron-root.log file is 0
    """
    log.info("Verifying Wireserver IP is reachable from root user")

    def check_exit_code(line):
        match = re.match("ExitCode:\\s(\\d+)", line)
        return match is not None and int(match.groups()[0]) == 0

    verify_data_in_cron_logs(cron_log=ROOT_CRON_LOG, verify=check_exit_code,
                             err_msg="Exit Code should be 0 for root based cron job!")

    if not check_wireserver_versions_file_exist(ROOT_WIRE_XML):
        fail("Wire version file should not be empty for root user!")


def verify_wireserver_ip_unreachable_for_non_root():
    """
    For non-root -
    Ensure the /tmp/wire-versions-non-root.xml is empty (generated by the cron job)
    Ensure the exit code in the /tmp/reboot-cron-non-root.log file is non-0
    """
    log.info("Verifying WireServer IP is unreachable from non-root user")

    def check_exit_code(line):
        match = re.match("ExitCode:\\s(\\d+)", line)
        return match is not None and int(match.groups()[0]) != 0

    verify_data_in_cron_logs(cron_log=NON_ROOT_CRON_LOG, verify=check_exit_code,
                             err_msg="Exit Code should be non-0 for non-root cron job!")

    if check_wireserver_versions_file_exist(NON_ROOT_WIRE_XML):
        fail("Wire version file should be empty for non-root user!")


def verify_tcp_connection_to_wireserver_for_non_root():
    """
    For non-root -
    Ensure the TCP 53 Connection exit code in the /tmp/reboot-cron-non-root.log file is 0
    """
    log.info("Verifying TCP connection to Wireserver port for non-root user")

    def check_exit_code(line):
        match = re.match("TCP 53 Connection ExitCode:\\s(\\d+)", line)
        return match is not None and int(match.groups()[0]) == 0

    verify_data_in_cron_logs(cron_log=NON_ROOT_CRON_LOG, verify=check_exit_code,
                             err_msg="TCP 53 Connection Exit Code should be 0 for non-root cron job!")


def generate_svg():
    """
    This is a good to have, but not must have. Not failing tests if we're unable to generate a SVG
    """
    log.info("Running systemd-analyze plot command to get the svg for boot execution order")
    dest_dir = os.path.join("/var", "log", "svgs")
    if not os.path.exists(dest_dir):
        os.makedirs(dest_dir)
    svg_name = os.path.join(dest_dir, "{0}.svg".format(BOOT_NAME))
    cmd = ["systemd-analyze plot > {0}".format(svg_name)]
    err_code, stdout = shellutil.run_get_output(cmd)
    if err_code != 0:
        log.info("Unable to generate svg: {0}".format(stdout))
    log.info("SVG generated successfully")


def main():
    try:
        # Verify firewall rules are set on boot through cron job logs
        verify_wireserver_ip_unreachable_for_non_root()
        verify_wireserver_ip_reachable_for_root()
        verify_tcp_connection_to_wireserver_for_non_root()
        verify_all_rules_exist()
    finally:
        # save the logs to /var/log to capture by collect-logs, this might be useful for debugging
        move_cron_logs_to_var_log()
        generate_svg()


parser = argparse.ArgumentParser()
parser.add_argument('-u', '--user', required=True, help="Non root user")
parser.add_argument('-bn', '--boot_name', required=True, help="Boot Name")
args = parser.parse_args()
NON_ROOT_USER = args.user
BOOT_NAME = args.boot_name
ROOT_CRON_LOG = "/tmp/reboot-cron-root.log"
NON_ROOT_CRON_LOG = f"/tmp/reboot-cron-{NON_ROOT_USER}.log"
NON_ROOT_WIRE_XML = f"/tmp/wire-versions-{NON_ROOT_USER}.xml"
ROOT_WIRE_XML = "/tmp/wire-versions-root.xml"
main()
