#!/usr/bin/env python3

import argparse
from datetime import datetime
import os
from datetime import timezone
from pathlib import Path

from jinja2 import Template

from sephiroth.providers import Provider
from sephiroth.providers.provider import supported_targets
import sephiroth

supported_servers = ["nginx", "apache", "caddy", "iptables", "ip6tables"]

base_dir = os.path.dirname(__file__)
template_dir = os.path.join(base_dir, "templates")


def get_output_path(args, build_date):
    """
    Input: Server type, date from build_template(), cloud provider
    Output: Path to file on disk to write to
    """
    targets_str = "_".join(args.targets)
    fdate = build_date.strftime("%Y-%m-%d_%H%M%S")
    fname = f"{fdate}_{args.servertype}_{targets_str}.sephiroth.conf"
    Path(args.output_dir).mkdir(exist_ok=True)
    return os.path.join(args.output_dir, fname)


def get_ranges(selected_provider, excludeip6=False, targets_in=None, compacted=False):
    """
    Input: Type of provider to target, as defined in supported_targets.
           Optionally exclude ip6, provide list of asns or files if asn or file target
    Output: Structured data ready to go to templates
    """
    if targets_in:
        provider = Provider(selected_provider, targets_in)
    else:
        provider = Provider(selected_provider)
    if compacted:
        return provider.get_compacted_ranges()
    return provider.get_processed_ranges()


def get_template(servertype):
    """
    Input: String name of server type
    Output: Jinja2 template object for given server type
    """
    fname = f"{servertype}/conf.jinja"
    template_path = os.path.join(template_dir, fname)
    return Template(open(template_path).read())


def build_template(ranges, template, build_date, use_proxy=False, redir_target=""):
    """
    Input: output of process_<provider>_ranges(), output of get_template()
    Output: Rendered template string ready to write to disk
    """
    return template.render(
        ranges=ranges["ranges"],
        header_comments=ranges["header_comments"],
        build_date=build_date,
        use_proxy=use_proxy,
        redir_target=redir_target,
    )


def print_output(servertype, targets, outfile):
    helpfile = os.path.join(template_dir, servertype, "help.jinja")
    abspath = os.path.abspath(outfile)
    targets_str = ", ".join(targets)
    help_text = Template(open(helpfile).read()).render(
        abspath=abspath, outfile=os.path.basename(outfile)
    )
    print(f"Your {servertype} blocklist for {targets_str} can be found at {outfile}\n")
    print(help_text)


def validate_nginx_args(args):
    if args.redir_target:
        print(
            "[?] Warning: We cannot generate nginx configs with redirect targets at this time. Ignoring."
        )
    return True


def validate_apache_args(args):
    if args.use_proxy:
        print(
            "[?] Warning: We cannot use PROXY protocol with Apache at this time. Ignoring."
        )
    if args.redir_target is None:
        print("[!] Error: Apache requires a defined redirect target using -r")
        raise SystemExit
    elif args.redir_target.startswith("http://") or args.redir_target.startswith(
        "https://"
    ):
        print(
            "[!] Error: Redirect target should not include scheme. Please edit the output RewriteRule directly if you want to change this."
        )
        raise SystemExit
    return True


def validate_caddy_args(args):
    if args.redir_target:
        print(
            "[?] Warning: We cannot generate caddy configs with redirect targets at this time. Ignoring."
        )
    return True


def validate_iptables_args(args):
    print("[?] Warning: iptables rules automatically exclude any IPv6 addresses")
    return True


def validate_ip6tables_args(args):
    print("[?] Warning: ip6tables rules automatically exclude any IPv4 addresses")
    return True


def parse_args():
    parser_desc = "Sephiroth is made to help block clouds."
    parser_epilog = "For more information, assistance, or to submit a pull request, please visit https://github.com/0xdade/sephiroth."
    parser = argparse.ArgumentParser(
        prog="Sephiroth", description=parser_desc, epilog=parser_epilog
    )
    parser.add_argument(
        "-s",
        "--server",
        help="Type of server to build blocklist for",
        required=True,
        choices=supported_servers,
        dest="servertype",
    )
    parser.add_argument(
        "-t",
        "--target",
        help="Targets to block",
        required=True,
        choices=supported_targets,
        action="append",
        dest="targets",
    )
    parser.add_argument(
        "-a",
        "--asn",
        help="ASN to block in AS#### format",
        action="append",
        metavar="ASN",
        dest="asns",
    )
    parser.add_argument(
        "-f",
        "--file",
        help="Files to block addresses from",
        action="append",
        metavar="FILENAME",
        dest="files",
    )
    parser.add_argument(
        "-r",
        "--redir",
        help="Place to redirect requests to. (apache)",
        default=None,
        dest="redir_target",
    )
    parser.add_argument(
        "-p",
        "--proxy",
        help="Using PROXY Protocol? (nginx)",
        default=False,
        action="store_true",
        dest="use_proxy",
    )
    parser.add_argument(
        "--no-ipv6",
        help="Exclude ipv6 addresses from the block list where applicable",
        default=False,
        action="store_true",
        dest="excludeip6",
    )
    parser.add_argument(
        "--compacted",
        help="Compact neighboring cidr ranges. This produces smaller file sizes but loses detail about each range.",
        default=False,
        action="store_true",
        dest="compacted",
    )
    parser.add_argument(
        "-o",
        "--output-dir",
        help="Specify an output dir",
        default=os.getcwd(),
        dest="output_dir",
    )
    parser.add_argument(
        "-V", "--version", action="version", version="%(prog)s " + sephiroth.__version__
    )
    return parser.parse_args()


def validate_targets(args):
    success = True
    for target in args.targets:
        if target == "asn" and not args.asns:
            print(
                "[!] Error: Cannot specify -t asn without including at least one -a AS####"
            )
            success = False
        elif target == "file" and not args.files:
            print(
                "[!] Error: Cannot specify -t file without including at least one -f filename.txt"
            )
            success = False
    return success


server_validators = {
    "apache": validate_apache_args,
    "nginx": validate_nginx_args,
    "caddy": validate_caddy_args,
    "iptables": validate_iptables_args,
    "ip6tables": validate_ip6tables_args,
}


def main():
    args = parse_args()
    if not validate_targets(args):
        raise SystemExit
    if args.servertype in server_validators:
        server_validators[args.servertype](args)
    build_date = datetime.now(timezone.utc)
    template_vars = {"header_comments": [], "ranges": []}
    if "_all" in args.targets:
        """
        _all can't be referenced as a provider because it's a meta target
        """
        excluded_from_all = {
            "_all",
        }
        targets = [t for t in supported_targets if t not in excluded_from_all]
    else:
        targets = args.targets
    for provider in targets:
        if args.asns and provider == "asn":
            provider_vars = get_ranges(
                provider, excludeip6=args.excludeip6, targets_in=args.asns
            )
        elif args.files and provider == "file":
            print(f"{args.files=}")
            provider_vars = get_ranges(
                provider, excludeip6=args.excludeip6, targets_in=args.files
            )
        elif provider in {"file", "asn"}:
            """
            If provider is file or asn but args.asns or args.files is empty, skip it
            """
            continue
        else:
            provider_vars = get_ranges(
                provider, excludeip6=args.excludeip6, compacted=args.compacted
            )
        template_vars["header_comments"] += provider_vars["header_comments"]
        template_vars["ranges"] += provider_vars["ranges"]
    template = get_template(args.servertype)
    template_output = build_template(
        template_vars, template, build_date, args.use_proxy, args.redir_target
    )
    outfile = get_output_path(args, build_date)
    with open(outfile, "w") as o:
        o.write(template_output)

    print_output(args.servertype, args.targets, outfile)


if __name__ == "__main__":
    main()
