#!/usr/bin/env python
"""
NetFoundry Tunnel Setup
"""
import sys
import os
import logging
import argparse
import semantic_version
import jwt
import reg_common


def check_registration(force):
    """
    Check if a successful registration happened previously.
    """
    if not force:
        if os.path.isfile("/opt/netfoundry/ziti/ziti-tunnel/.is_registered"):
            print(
                "This tunnel has been previously registered. "
                "If you wish to override the current registration, add -f/--force to the command")
            sys.exit(1)
    reg_common.manage_services("ziti-tunnel","stop")

def identify_os():
    """
    Try to identify the OS so we know how to change
    """
    if os.path.isfile("/etc/lsb-release"):
        with open("/etc/lsb-release", 'r') as file:
            release_content = file.read()
            if "20.04" in release_content:
                os_version = "Ubuntu20"
    elif os.path.isfile("/etc/os-release"):
        with open("/etc/os-release", 'r') as file:
            release_content = file.read()
            if "Raspbian" in release_content:
                os_version = "Raspbian"
    else:
        print("ERROR: Un-supported OS")
        print("Supported OS: Ubuntu20.04, Raspbian")
        sys.exit(1)

    return os_version

def run_registration(args, jwt_path):
    """
    Run the registration logic
    """
    # identify the OS
    identify_os()

    # check for previous registrations
    check_registration(args.force)

    print("\033[0;35mStarting Registration\033[0m")

    # check jwt file
    logging.debug("Attempting to read jwt file")
    print("Checking jwt")
    try:
        if os.path.isfile(jwt_path):
            with open(jwt_path, 'r') as file:
                jwt_encoded_content = file.read().replace('\n', '')
        else:
            print("ERROR: Unable to find jwt file")
            sys.exit(1)
    except OSError:
        print("ERROR accessing jwt file")
    # extract jwt file
    try:
        jwt_decoded_content = jwt.decode(jwt_encoded_content, algorithms=['RS256'], verify=False)
    except jwt.DecodeError:
        print("ERROR: Unable to parse jwt")
        sys.exit(1)
    logging.debug('JWT Content: %s', jwt_decoded_content)

    # extract controller url from jwt content
    controller_url = jwt_decoded_content['iss']

    # get ziti version from controller
    print("Contacting Controller for version information")
    ziti_version = reg_common.get_ziti_controller_version(controller_url)

    print("Starting ziti setup")
    # create directories and download binaries
    ziti_tunnel_setup(ziti_version, jwt_path)

    if not args.host:
        # interface check
        adapter_ip, adapter_name, adapter_prefix = reg_common.tunnler_adapter_check(args.tunnel_ip)

        # setup local dns resolution
        print("Starting dns setup")
        reg_common.dns_setup(adapter_ip)
        reg_common.ufw_add_rules(adapter_ip,adapter_prefix,"53","udp","local")
        reg_common.ufw_add_rules(adapter_ip,adapter_prefix,"53","tcp","local")
        ziti_service_files(ziti_version, adapter_ip, adapter_name)
    else:
        # create service files
        ziti_service_files(ziti_version)

    # start services
    reg_common.manage_services("ziti-tunnel","start")

    print("\033[0;35mRegistration complete\033[0m")

def ziti_tunnel_setup(version, jwt_path):
    """
    Download Ziti Tunnel
    """
    # create dirs if they don't exist
    logging.debug("Checking if Ziti directories exist")
    try:
        reg_common.create_directory(directory_path="/var/log/ziti")
        reg_common.create_directory(directory_path="/opt/netfoundry/ziti")
        reg_common.create_directory(directory_path="/opt/netfoundry/ziti/ziti-tunnel")
    except OSError:
        print("ERROR: Unable to create ziti directories")
        sys.exit(1)

    # set architecture
    arm_machine = reg_common.arm_check()

    product_metadata = reg_common.get_mop_product_metadata(version)


    download_url = reg_common.parse_metadata(product_metadata,
                                             arm_machine,
                                             "zitiBinaryBundleLinuxAMD64")

    # download binaries
    reg_common.download_bundle(download_url)

    # run enrollment
    ziti_register(jwt_path)

def ziti_register(jwt_path):
    """
    Register the ziti tunnel
    This function should be updated once Ziti has fix the exit codes.
    """
    command_output = reg_common.run_cmd(
        ["/opt/netfoundry/ziti/ziti-tunnel/ziti-tunnel",
         "enroll",
         "-j",
         jwt_path,
         "-o",
         "config.json"], "/opt/netfoundry/ziti/ziti-tunnel")

    if "enrolled successfully" in command_output.stderr:
        logging.debug("Successfully registed Ziti")
        try:
            open("/opt/netfoundry/ziti/ziti-tunnel/.is_registered", "w").close()
        except OSError:
            logging.debug("\033[0;31mERROR:\033[0m writing jwt file to disk")
    else:
        print("\033[0;31mERROR:\033[0m Failed to registed Ziti")
        sys.exit(1)

def ziti_service_files(ziti_version, adapter_ip=None, adapter_name=None):
    """
    Create Ziti Tunnel service file
    """
    print("Installing Systemd Service")
    #Can be removed once all networks are 0.18.9 and above.
    if adapter_ip is None:
        logging.debug("Writing service file in host only mode")
        service_content = ("[Unit]\n"
                        "Description=Ziti-Tunnel\n"
                        "After=network.target\n"
                        "\n"
                        "[Service]\n"
                        "User=root\n"
                        "WorkingDirectory=/opt/netfoundry/ziti/ziti-tunnel\n"
                        "ExecStart=/opt/netfoundry/ziti/ziti-tunnel/ziti-tunnel"
                        " host -i /opt/netfoundry/ziti/ziti-tunnel/config.json\n"
                        "Restart=always\n"
                        "RestartSec=2\n"
                        "LimitNOFILE=65535\n"
                        "\n"
                        "[Install]\n"
                        "WantedBy=multi-user.target\n")
    else:
        if semantic_version.Version(ziti_version) > semantic_version.Version("0.18.8"):
            logging.debug("Writing service file with --lanIf")
            service_content = ("[Unit]\n"
                            "Description=Ziti-Tunnel\n"
                            "After=network.target\n"
                            "\n"
                            "[Service]\n"
                            "User=root\n"
                            "WorkingDirectory=/opt/netfoundry/ziti/ziti-tunnel\n"
                            "ExecStart=/opt/netfoundry/ziti/ziti-tunnel/ziti-tunnel"
                            " tproxy -i /opt/netfoundry/ziti/ziti-tunnel/config.json"
                            " -r udp://" + adapter_ip + ":53 --lanIf " + adapter_name + "\n"
                            "Restart=always\n"
                            "RestartSec=2\n"
                            "LimitNOFILE=65535\n"
                            "\n"
                            "[Install]\n"
                            "WantedBy=multi-user.target\n")
        else:
            logging.debug("Writing service file without --lanIf")
            service_content = ("[Unit]\n"
                            "Description=Ziti-Tunnel\n"
                            "After=network.target\n"
                            "\n"
                            "[Service]\n"
                            "User=root\n"
                            "WorkingDirectory=/opt/netfoundry/ziti/ziti-tunnel\n"
                            "ExecStart=/opt/netfoundry/ziti/ziti-tunnel/ziti-tunnel"
                            " run /opt/netfoundry/ziti/ziti-tunnel/config.json"
                            " -r udp://" + adapter_ip + ":53\n"
                            "Restart=always\n"
                            "RestartSec=2\n"
                            "LimitNOFILE=65535\n"
                            "\n"
                            "[Install]\n"
                            "WantedBy=multi-user.target\n")

    # write the service file
    try:
        reg_common.create_file(name="ziti-tunnel.service",
                               path="/etc/systemd/system/",
                               content=service_content)
    except OSError:
        print("\033[0;31mERROR:\033[0m Unable to write service file")
        sys.exit(1)

def main():
    """
    Main logic
    """
    __version__ = '4.1.0'
    # change log
    # 1.0.0 - initial release
    # 1.1.0 - accept jwt string or file
    # 1.1.1 - changed service file to output pfxlog
    # 1.1.2 - fixed bug with jwt file
    # 1.1.3 - set locale to en_US
    # 1.2.0 - adjust system service file to force nofile to 65k
    # 1.2.1 - fixed PyJWT to version 1.7.1 in requirements file
    # 1.2.2 - reverted pfxlog to json; added log output to systemd
    # 1.2.3 - add raspbian as supported os
    # 2.0.0 - drop centos support; switch tunnel to tun mode;
    # enforce forwarding on OS; add UFW rules
    # 2.0.1 - specified mtu size of 9000 in service file
    # 2.1.0 - reverted to tproxy, removed forwarding enforcement &
    # fowarding firewall rules; drop mtu
    # 2.2.0 - added download of cli if it's not already there
    # 2.2.1 - revert back to journald for logging
    # 2.2.2 - print out when writing service file
    # 2.3.0 - add logic to handle the lanIf option for version 0.18.9 and higher
    # 2.3.1 - fix local issue by forcing write with utf-8
    # 3.0.0 - get download url from MOP
    # 3.1.0 - split reg_common
    # 3.2.0 - add host only mode
    # 3.2.1 - enforce permissions on created files
    # 4.0.0 - switch from product-metadata to network-version endpoint
    # 4.1.0 - change download to bundle

    # argument parser
    parser = argparse.ArgumentParser()
    # arguments
    jwt_group = parser.add_mutually_exclusive_group(required=True)
    jwt_group.add_argument('jwt_string', nargs='?',
                           help='JWT String', default=None)
    jwt_group.add_argument('-j', '--jwt', dest='identity_jwt_path', nargs='?',
                           help="Path to JWT file", default=None)
    config_group =  parser.add_mutually_exclusive_group(required=False)
    config_group.add_argument('-o', '--tunnel_ip',
                              help="Override the IP Address used to service DNS queries")
    config_group.add_argument('--host',
                              action="store_true",
                              help="Configure for host only mode & skip DNS configuration")
    parser.add_argument('-f', '--force',
                        action="store_true",
                        help='Forcefully proceed with re-registration')
    parser.add_argument('-d', '--debug',
                        action="store_true",
                        help='Enable debug log in log file output')
    parser.add_argument('-v', '--version',
                        action='version',
                        version=__version__)

    # get arguments passed
    args = parser.parse_args()

    # enable debug if requested
    if args.debug:
        log_level = logging.DEBUG
    else:
        log_level = logging.INFO

    # define logging information
    logging.basicConfig(format='%(asctime)s-%(levelname)s-%(message)s',
                        datefmt='%Y-%m-%d-%H:%M:%S',
                        level=log_level
                        )

    # run root check
    reg_common.root_check()

    # systemd check
    reg_common.systemd_check()

    # handle jwt format
    if args.identity_jwt_path is not None:
        jwt_path = os.path.realpath(args.identity_jwt_path)
    elif args.jwt_string is not None:
        try:
            if os.path.isfile("/tmp/jwt"):
                os.remove("/tmp/jwt")
        except OSError:
            print("\033[0;31mERROR:\033[0m Unable to remove previous /tmp/jwt file")
        try:
            with open("/tmp/jwt", 'w', encoding="utf-8") as file:
                file.write(args.jwt_string)
                jwt_path = "/tmp/jwt"
        except OSError:
            print("\033[0;31mERROR:\033[0m: Unable to write /tmp/jwt file")
            sys.exit(1)
    else:
        parser.print_help()

    if not args.host:
        # ufw check
        reg_common.ufw_check()

        # adapter/ip count check
        if not args.tunnel_ip:
            reg_common.adapter_count_check("tunnel")

        # check override is actually an IP Address
        if args.tunnel_ip:
            logging.debug("Checking if override is valid")
            override_is_ip = reg_common.ip_check(args.tunnel_ip)
            if not override_is_ip:
                print("\033[0;31mERROR:\033[0m The override IP was not a valid IP Address")
                sys.exit(1)


    # run registration
    run_registration(args, jwt_path)

# main
if __name__ == '__main__':
    main()