# Example usage:
# sudo python3 sshd_running_config_nm.py --ScriptName sshd_running_config_nm.py --CollectConfig <comma-separated list of config keys>
# sudo python2 sshd_running_config_nm.py --ScriptName sshd_running_config_nm.py --CollectConfig <comma-separated list of config keys>

import os
import sys
import subprocess
import argparse
import json
import traceback

class FatalArgumentError(Exception):
    pass

def write_stderr(msg):
    try:
        sys.stderr.write(msg)
        sys.stderr.flush()
    except Exception:
        pass

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--ScriptName", dest='script_name', help='name of the script')
    parser.add_argument("--CollectConfig", dest='collect_config_list',  required=True,
                        help='comma-separated list of config to collect')
    args = parser.parse_args()
    return args

def run_sshd_T(user):
    try:
        cmd = [
            '/usr/sbin/sshd', '-T', '-C', 'user={}'.format(user)
        ]
        # Python 2/3 compatible subprocess
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = proc.communicate()
        # In Python 3, stdout/stderr are bytes, decode to str
        if sys.version_info[0] >= 3:
            if isinstance(stdout, bytes):
                stdout = stdout.decode('utf-8', 'replace')
            if isinstance(stderr, bytes):
                stderr = stderr.decode('utf-8', 'replace')
        if proc.returncode != 0:
            write_stderr("sshd -T failed: {}\n".format(stderr.strip()))
            raise RuntimeError("sshd -T failed")
        return stdout
    except OSError:
        write_stderr("Error: sshd command not found.\n")
        raise
    except Exception as e:
        write_stderr("Exception in fetching the config: {}\n".format(e))
        raise

def parse_sshd_output_to_json(sshd_output):
    config = {}
    for line in sshd_output.splitlines():
        if not line.strip():
            continue
        if ' ' in line:
            key, value = line.split(None, 1)
            if key in config:
                if isinstance(config[key], list):
                    config[key].append(value)
                else:
                    config[key] = [config[key], value]
            else:
                config[key] = value
    return config

def main(args):
    if os.geteuid() != 0:
        raise ValueError("script needs to be run as root")

    try:
        sshd_output = run_sshd_T('root')
        config_json = parse_sshd_output_to_json(sshd_output)

        # Prepare the list of keys to collect, case-insensitive
        collect_keys = [k.strip().lower() for k in args.collect_config_list.split(',') if k.strip()]
        # convert the json to lower case keys for case-insensitive comparison
        config_json_lower = {k.lower(): v for k, v in config_json.items()}
        # Filter the config_json to only include keys that are in collect_keys
        filtered_json = {k: config_json_lower[k] for k in collect_keys if k in config_json_lower}
        keys_not_found = [k for k in collect_keys if k not in config_json_lower]

        payload = dict(scriptVersion=1,
                cmdline=" ".join(sys.argv[1:]),
                missingConfig=keys_not_found,
                collectedConfig=filtered_json)

        payload_json = json.dumps(payload)

        # write to stderr since stdout is scrubbed out and can't be used for analysis.
        write_stderr(payload_json + "\n")
        return 0

    except Exception:
        write_stderr(traceback.format_exc())
        return 3

if __name__ == "__main__":
    try:
        args = parse_args()
        sys.exit(main(args))
    except FatalArgumentError:
        sys.exit(2)
    except SystemExit as e:
        sys.exit(getattr(e, "code", 1))
    except Exception:
        write_stderr(traceback.format_exc())
        sys.exit(3)