#!/usr/bin/python

# Copyright 2016 Red Hat Inc., Durham, North Carolina.
# All Rights Reserved.
#
# openscap-daemon is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 2.1 of the License, or
# (at your option) any later version.
#
# openscap-daemon is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.

# You should have received a copy of the GNU Lesser General Public License
# along with openscap-daemon.  If not, see <http://www.gnu.org/licenses/>.
#
# Authors:
#   Martin Preisler <mpreisle@redhat.com>

from openscap_daemon import config as config_
from openscap_daemon import evaluation_spec
from openscap_daemon import oscap_helpers
from openscap_daemon import cli_helpers
from openscap_daemon import version
from openscap_daemon.evaluation_spec import ProfileSuffixMatchError

import os
import os.path
import logging
import argparse
import sys
import threading
import io
import json
import datetime

if sys.version_info < (3,):
    import Queue
else:
    import queue as Queue


def cli_xml(args, config):
    spec = evaluation_spec.EvaluationSpec()
    spec.load_from_xml_file(args.path)
    results, stdout, stderr, exit_code = spec.evaluate(config)
    if args.results is not None:
        args.results.write(results)
        args.results.close()
    if args.stdout is not None:
        args.stdout.write(stdout)
        args.stdout.close()
    if args.stderr is not None:
        args.stderr.write(stderr)
        args.stderr.close()

    sys.exit(exit_code)


def cli_spec(args, config):
    spec = evaluation_spec.EvaluationSpec()
    spec.mode = oscap_helpers.EvaluationMode.from_string(args.mode)
    spec.target = args.target
    if spec.mode not in [oscap_helpers.EvaluationMode.CVE_SCAN,
                         oscap_helpers.EvaluationMode.STANDARD_SCAN]:
        spec.input_.set_contents(args.input_.read())
    if args.tailoring is not None:
        spec.tailoring.set_contents(args.tailoring.read())

    spec.profile_id = args.profile
    spec.online_remediation = args.remediate

    if args.print_xml:
        print(spec.to_xml_source())
        sys.exit(0)

    else:
        results, stdout, stderr, exit_code = spec.evaluate(config)
        if args.results is not None:
            args.results.write(results)
            args.results.close()
        if args.stdout is not None:
            args.stdout.write(stdout)
            args.stdout.close()
        if args.stderr is not None:
            args.stderr.write(stderr)
            args.stderr.close()

        sys.exit(exit_code)


def cli_scan(args, config):
    assert(os.path.isdir(args.output))

    output_dir_map = {}
    targets = cli_helpers.preprocess_targets(args.targets, output_dir_map)

    queue = Queue.Queue(len(targets))
    for target in targets:
        queue.put_nowait(target)

    scanned_targets = []
    failed_targets = []

    def scan_worker():
        while True:
            try:
                target = queue.get(False)

                if len(failed_targets) > 0:
                    failed_targets.append(target)
                    queue.task_done()
                    continue

                cve_results = None
                cve_last_updated = None
                standard_scan_results = None

                logging.debug("Started scanning target '%s'", target)
                started_time = None
                finished_time = None

                try:
                    started_time = datetime.datetime.now()
                    cpes = []
                    try:
                        cpes = evaluation_spec.EvaluationSpec.\
                            detect_CPEs_of_target(target, config)
                    except:
                        logging.exception(
                            "Failed to detect CPEs of target '%s'. "
                            "Assuming no CPEs..." % (target)
                        )

                    if not args.no_cve_scan:
                        es = evaluation_spec.EvaluationSpec()
                        es.mode = oscap_helpers.EvaluationMode.CVE_SCAN
                        es.target = target
                        es.cpe_hints = cpes
                        try:
                            cve_results, stdout, stderr, exit_code = \
                                es.evaluate(config)

                            if exit_code == 1:
                                logging.warning(
                                    "CVE scan of target '%s' failed with "
                                    "exit_code %i.\n\nstdout:%s\n\nstderr:%s" %
                                    (target, exit_code, stdout, stderr)
                                )
                        except:
                            logging.exception(
                                "Failed to scan target '%s' for "
                                "vulnerabilities." % (target)
                            )

                        try:
                            cve_last_updated = config.cve_feed_manager.\
                                get_cve_feed_last_updated(cpes)
                        except:
                            # this is not a crucial part of evaluation, the
                            # last modified date can be unknown.
                            pass

                    if not args.no_standard_compliance:
                        es = evaluation_spec.EvaluationSpec()
                        es.mode = oscap_helpers.EvaluationMode.STANDARD_SCAN
                        es.target = target
                        es.cpe_hints = cpes
                        ssg_sds = config.get_ssg_sds(cpes)
                        es.input_.set_file_path(ssg_sds)
                        try:
                            args.profile = es.select_profile_by_suffix(args.profile)
                        except ProfileSuffixMatchError as e:
                            logging.error(
                                "Failed to find profile matching suffix '%s' "
                                "in profiles applicable to target '%s': %s"
                                % (args.profile, target, e.message)
                            )
                            raise RuntimeError
                        try:
                            standard_scan_results, stdout, stderr, exit_code = \
                                es.evaluate(config)

                            if exit_code == 1:
                                logging.warning(
                                    "Configuration compliance scan of target '%s' "
                                    "using profile '%s' "
                                    "failed with exit_code %i.\n\nstdout:%s\n\n"
                                    "stderr:%s" %
                                    (target, es.profile_id, exit_code, stdout, stderr)
                                )

                        except:
                            logging.exception(
                                "Failed to scan target '%s' for "
                                "configuration compliance." % (target)
                            )

                    finished_time = datetime.datetime.now()

                except Exception as e:
                    logging.exception(e)
                    failed_targets.append(target)

                queue.task_done()
                scanned_targets.append(
                    (target, cve_results, cve_last_updated,
                     standard_scan_results, started_time, finished_time)
                )

                percent = "{0:6.2f}%".format(
                    float(len(scanned_targets) * 100) / len(targets)
                )

                logging.info("[%s] Scanned target '%s'", percent, target)

            except Queue.Empty:
                break

    assert(args.jobs > 0)

    workers = []
    for worker_id in range(args.jobs):
        worker = threading.Thread(
            name="Atomic scan worker #%i" % (worker_id),
            target=scan_worker
        )
        workers.append(worker)
        worker.start()

    try:
        queue.join()

    except KeyboardInterrupt:
        failed_targets.append(None)

        for worker in workers:
            worker.join()

        sys.stderr.write("Evaluation interrupted by user!\n")

    if len(failed_targets) > 0:
        sys.stderr.write(
            "Fatal error encountered while evaluating! Failed to evaluate "
            "at least %i targets!\n" % (len(failed_targets) - 1)
        )

    for target, cve_results, cve_last_updated, standard_scan_results, \
            started_time, finished_time in scanned_targets:
        output_dir = ""
        if target in output_dir_map:
            output_dir = output_dir_map[target]

        else:
            output_dir = target
            output_dir = output_dir.replace(":", "_")
            output_dir = output_dir.replace("/", "_")

        json_target = target
        if json_target.startswith("chroot://"):
            json_target = json_target[len("chroot://"):]

        json_data = {}
        json_data["UUID"] = json_target
        json_data["Scanner"] = "openscap"
        json_data["Time"] = started_time.strftime("%Y-%m-%dT%H:%M:%S") \
            if started_time is not None else "unknown"
        json_data["Finished Time"] = \
            finished_time.strftime("%Y-%m-%dT%H:%M:%S") \
            if finished_time is not None else "unknown"
        if cve_results is not None:
            json_data["CVE Feed Last Updated"] = \
                cve_last_updated.strftime("%Y-%m-%dT%H:%M:%S") \
                if cve_last_updated is not None else "unknown"
        json_data["Vulnerabilities"] = []
        if (args.no_cve_scan or cve_results) and \
                (args.no_standard_compliance or standard_scan_results):
            json_data["Successful"] = "true"
        else:
            json_data["Successful"] = "false"

        scan_type = []
        full_output_dir = os.path.join(args.output, output_dir)
        try:
            os.makedirs(full_output_dir)
        except OSError as e:
            if e.errno != 17:  # it's fine if it already exists
                raise

        if cve_results is not None:
            scan_type.append("CVE")

            cli_helpers.summarize_cve_results(
                cve_results, json_data["Vulnerabilities"]
            )

            with io.open(os.path.join(
                    full_output_dir, "cve.xml"), "w",
                    encoding="utf-8") as f:
                f.write(cve_results)

        if standard_scan_results is not None:
            scan_type.append("Configuration Compliance")

            cli_helpers.summarize_standard_compliance_results(
                standard_scan_results, json_data["Vulnerabilities"], args.profile
            )
            json_data["Profile"] = args.profile

            arf_filepath = os.path.join(full_output_dir, "arf.xml")
            with io.open(arf_filepath, "w", encoding="utf-8") as f:
                f.write(standard_scan_results)
            if args.fix_type is not None:
                fix_script = oscap_helpers.generate_fix_for_result(config, arf_filepath, args.fix_type)
                suffixes = {"bash": "sh", "ansible": "yml", "puppet": "pp"}
                fix_name = "fix." + suffixes[args.fix_type]
                fix_filepath = os.path.join(full_output_dir, fix_name)
                with io.open(fix_filepath, "w", encoding="utf-8") as f:
                    f.write(fix_script)
            if args.report:
                report = oscap_helpers.generate_html_report_for_result(config, arf_filepath)
                report_filepath = os.path.join(full_output_dir, "report.html")
                with io.open(report_filepath, "w", encoding="utf-8") as f:
                    f.write(report)

        json_data["Scan Type"] = ", ".join(scan_type)

        with open(os.path.join(
                full_output_dir, "json"), "w") as f:
            json.dump(json_data, f, indent=2)


def main():
    parser = argparse.ArgumentParser(
        description="OpenSCAP-Daemon one-off evaluator."
    )
    parser.add_argument(
        "-v", "--version", action="version",
        version="%(prog)s " + version.VERSION_STRING
    )
    parser.add_argument("--verbose",
                        help="be verbose, useful for debugging",
                        action="store_true")

    subparsers = parser.add_subparsers(dest="action")
    subparsers.required = True

    config_parser = subparsers.add_parser(
        "config",
        help="Start with default configuration, auto-detect tool and content "
             "locations and output the resulting INI results into stdout or "
             "given file path"
    )
    config_parser.add_argument(
        "--path", metavar="PATH", type=argparse.FileType("w"),
        default=sys.stdout,
        help="Destination where the config file will be written, defaults to "
             "stdout."
    )

    xml_parser = subparsers.add_parser(
        "xml",
        help="Evaluate an EvaluationSpec passed as an XML, either to stdin or "
             "as a file."
    )
    xml_parser.add_argument(
        "--path", metavar="PATH", type=argparse.FileType("r"),
        default=sys.stdin,
        help="The input Evaluation Spec XML file. Defaults to stdin."
    )
    xml_parser.add_argument(
        "--results", metavar="PATH", type=argparse.FileType("w"),
        help="Write ARF (result datastream) or OVAL results XML to this file."
    )
    xml_parser.add_argument(
        "--stdout", metavar="PATH", type=argparse.FileType("w"),
        help="Write stdout from oscap tool to this file."
    )
    xml_parser.add_argument(
        "--stderr", metavar="PATH", type=argparse.FileType("w"),
        help="Write stderr from oscap tool to this file."
    )

    spec_parser = subparsers.add_parser(
        "spec",
        help="Evaluate an EvaluationSpec created using arguments passed on "
             "the command line."
    )
    spec_parser.add_argument(
        "--mode", type=str,
        choices=["sds", "oval", "cve_scan", "standard_scan"],
        default="sds",
        help="Evaluation mode for the EvaluationSpec. 'sds' evaluates input as "
        "a source datastream. 'oval' evaluates it as an OVAL file. 'cve_scan' "
        "is a special mode that automatically uses the right CVE feed as OVAL "
        "file. 'standard_scan' uses the right SSG content and standard profile "
        "based on OS of the scanned system."
    )
    spec_parser.add_argument(
        "--target", type=str,
        default="localhost",
        help="Which target should we be evaluating. Possible choices include: "
        "'localhost', 'ssh://user@machine', 'docker-image://IMAGE_ID', "
        "'docker-container://CONTAINER_ID', 'vm-domain://VM_NAME', "
        "'vm-image:///path/to/image.qcow2', 'chroot:///path/to/chroot'."
    )
    spec_parser.add_argument(
        "--input", metavar="PATH", dest="input_",
        type=lambda path: io.open(path, "r", encoding="utf-8"),
        default=sys.stdin,
        help="Depending on --mode this should be a source datastream or OVAL "
        "file. In cve_scan and standard_scan mode the --input is not used."
    )
    spec_parser.add_argument(
        "--tailoring", metavar="PATH",
        type=lambda path: io.open(path, "r", encoding="utf-8"),
        help="XCCDF tailoring file. Only used in 'sds' mode."
    )
    spec_parser.add_argument(
        "--profile", type=str,
        default="",
        help="ID of the XCCDF profile to use. Only used in 'sds' mode. Empty "
        "string is the default and that means the (default) profile."
    )
    spec_parser.add_argument(
        "--remediate", default=False, action="store_true",
        help="Perform remediation for failed rules after the scan. Only used "
        "in 'sds' and 'standard_scan' modes."
    )
    spec_parser.add_argument(
        "--print-xml",
        dest="print_xml",
        action="store_true",
        help="Don't evaluate the EvaluationSpec, just print its XML to stdout"
    )
    spec_parser.add_argument(
        "--results", metavar="PATH",
        type=lambda path: io.open(path, "w", encoding="utf-8"),
        help="Write OVAL results or ARF result datastream (depending on mode) "
        "to this location."
    )
    spec_parser.add_argument(
        "--stdout", metavar="PATH",
        type=lambda path: io.open(path, "w", encoding="utf-8"),
        help="Write stdout from oscap tool to this file."
    )
    spec_parser.add_argument(
        "--stderr", metavar="PATH",
        type=lambda path: io.open(path, "w", encoding="utf-8"),
        help="Write stderr from oscap tool to this file."
    )
    target_cpes_parser = subparsers.add_parser(
        "target-cpes",
        help="Detect CPEs applicable on given target"
    )
    target_cpes_parser.add_argument(
        "--target", type=str,
        default="localhost",
        help="Which target should we be checking. Possible choices include: "
        "'localhost', 'ssh://user@machine', 'docker-image://IMAGE_ID', "
        "'docker-container://CONTAINER_ID', 'vm-domain://VM_NAME', "
        "'vm-image:///path/to/image.qcow2', 'chroot:///path/to/chroot'."
    )
    target_profiles_parser = subparsers.add_parser(
        "target-profiles",
        help="Detect SCAP Security Guide profiles applicable on given target"
    )
    target_profiles_parser.add_argument(
        "--target", type=str,
        default="localhost",
        help="Which target should we be checking. Possible choices include: "
             "'localhost', 'ssh://user@machine', 'docker-image://IMAGE_ID', "
             "'docker-container://CONTAINER_ID', 'vm-domain://VM_NAME', "
             "'vm-image:///path/to/image.qcow2', 'chroot:///path/to/chroot'."
    )
    scan_parser = subparsers.add_parser(
        "scan",
        help="Scan a list of targets for CVEs and configuration compliance, "
        "return aggregated results. This is an integration shim "
        "intended for Atomic but can also be useful elsewhere."
    )
    scan_parser.add_argument(
        "--targets", type=str, nargs="+",
        default=["localhost"],
        help="Which target(s) should we be scanning. Possible choices include: "
        "'localhost', 'ssh://user@machine', 'docker-image://IMAGE_ID', "
        "'docker-container://CONTAINER_ID', 'vm-domain://VM_NAME', "
        "'vm-image:///path/to/image.qcow2', 'chroot:///path/to/chroot', "
        "'chroots-in-dir:///path/to/chroots'. "
        "Delimited by spaces."
    )
    scan_parser.add_argument(
        "-j", "--jobs", type=int,
        default=4,
        help="How many worker jobs should scan in parallel."
    )
    scan_parser.add_argument(
        "--no-cve-scan", default=False, action="store_true",
        dest="no_cve_scan",
        help="Skip the CVE scan."
    )
    scan_parser.add_argument(
        # keeping the alias for compatibility with Atomic
        "--no-configuration-compliance", "--no-standard-compliance", default=False, action="store_true",
        dest="no_standard_compliance",
        help="Skip the configuration compliance scan."
    )
    scan_parser.add_argument(
        "--profile", type=str,
        default="xccdf_org.ssgproject.content_profile_standard",
        help="Specify the profile ID for configuration compliance scan. "
        "If not specified, the 'standard' profile will be used."
    )
    scan_parser.add_argument(
        "--output", type=str, required=True,
        help="A directory where results will be stored in. There will be a "
        "directory for each target created there with up to 4 files. 'json' "
        "with json summary of the scan, cve.xml with CVE scan raw results, "
        "arf.xml with configuration compliance scan raw results, and "
        "fix.[sh|yml|pp] with a compliance remediation script."
    )
    scan_parser.add_argument(
        "--fix_type", type=str,
        choices=["bash", "ansible", "puppet"], default=None,
        help="Specify the language of remediation script to be used."
    )
    scan_parser.add_argument(
        "--report", action="store_true", default=False,
        help="Create HTML report in the output directory."
    )
    args = parser.parse_args()

    logging.basicConfig(format='%(levelname)s:%(message)s',
                        level=logging.DEBUG if args.verbose else logging.INFO)
    logging.info("OpenSCAP Daemon one-off evaluator %s", version.VERSION_STRING)

    if args.action == "config":
        config = config_.Configuration()
        config.autodetect_tool_paths()
        config.autodetect_content_paths()
        config.save_as(args.path)
        sys.exit(0)

    config_file = os.path.join("/", "etc", "oscapd", "config.ini")
    if "OSCAPD_CONFIG_FILE" in os.environ:
        config_file = os.environ["OSCAPD_CONFIG_FILE"]

    config = config_.Configuration()
    config.load(config_file)
    config.autodetect_tool_paths()
    config.autodetect_content_paths()
    config.prepare_dirs(cleanup_allowed=False)
    try:
        config.sanity_check()
    except:
        logging.exception("Configuration file failed sanity checking!")
        sys.exit(1)

    if args.action == "xml":
        cli_xml(args, config)

    elif args.action == "spec":
        cli_spec(args, config)

    elif args.action == "target-cpes":
        cpes = evaluation_spec.EvaluationSpec.detect_CPEs_of_target(
            args.target, config
        )
        print("\n".join(cpes))
        sys.exit(0)

    elif args.action == "target-profiles":
        cpes = evaluation_spec.EvaluationSpec.detect_CPEs_of_target(
            args.target, config
        )
        ssg_sds = config.get_ssg_sds(cpes)
        print("Security profiles applicable on target " + args.target + ":")
        profiles = oscap_helpers.get_profile_choices_for_input(ssg_sds, None)
        for profile_id, title in profiles.items():
            if profile_id:
                print(title + " (id='" + profile_id + "')")
            else:
                print("Default Profile")
        sys.exit(0)

    elif args.action == "scan":
        cli_scan(args, config)


if __name__ == "__main__":
    main()
