#!/usr/bin/env python
#
# Copyright (c) 2018-2020  Joe Clarke <jclarke@cisco.com>
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.


from __future__ import print_function
from meraki_api import Meraki, Network, Vlan, SwitchPort, SSID, Device
import yaml
import argparse
import sys
import os
from colorama import Fore, Style
import colorama

try:
    from yaml import CLoader as Loader
except ImportError:
    from yaml import Loader

BANNER = "[{}] **********************************************************"


def main():
    # Setup command line arguments.
    parser = argparse.ArgumentParser(prog=sys.argv[0], description="Add devices to network")
    parser.add_argument("--config", "-c", metavar="<CONFIG FILE>", help="Path to the organization configuration file", required=True)
    parser.add_argument("--networks", "-n", metavar="<NETWORK>[,<NETWORK>[,...]]", help="Comma-separated list of networks to process")
    args = parser.parse_args()

    colorama.init()

    if not os.path.isfile(args.config):
        print("Config file {} does not exist or is not a file!".format(args.config))
        exit(1)

    print(BANNER.format("Loading config file"))
    with open(args.config, "r") as c:
        config = yaml.load(c, Loader=Loader)
    print("{}ok{}\n".format(Fore.GREEN, Style.RESET_ALL))

    for key in ["api_key", "organization", "networks"]:
        if key not in config:
            print("Invalid config: {} is missing!".format(key))
            exit(1)

    meraki = Meraki(key=config["api_key"])
    orgs = meraki.get_organizations()
    org = None
    for o in orgs:
        if o.get("name") == config["organization"]:
            org = o
            break

    if org is None:
        print("Failed to find organization {} in this profile!".format(config["organization"]))
        exit(1)

    nets = org.get_networks()
    inv = org.get_inventory()

    errors = 0

    configure_nets = None
    if args.networks is not None:
        configure_nets = args.networks.split(",")

    for nname, network in config["networks"].items():
        nerrors = 0
        net_obj = None
        print(BANNER.format("Configuring network {}".format(nname)))
        if configure_nets is not None and nname not in configure_nets:
            print("{}skipping (not in specified network list){}".format(Fore.BLUE, Style.RESET_ALL))
            continue

        validn = True
        for key in ["address", "timezone"]:
            if key not in network:
                print("{}Invalid network config for {}: {} is missing!{}".format(Fore.RED, nname, key, Style.RESET_ALL))
                errors += 1
                validn = False
                break

        if not validn:
            continue

        for n in nets:
            if n.get("name") == nname:
                net_obj = n
                break

        if net_obj is None:
            nargs = {"timezone": network["timezone"]}
            if "copy_from_network" in network:
                for n in nets:
                    if n.get("name") == network["copy_from_network"]:
                        nargs["copy_from_network_id"] = n.get("id")
                        break
            net_obj = org.create_network(nname, **nargs)

        if net_obj is None:
            print("{}Error creating new network {}!{}".format(Fore.RED, nname, Style.RESET_ALL))
            errors += 1
            continue

        if "client_bandwidth_limit" in network:
            if not isinstance(network["client_bandwidth_limit"], int) or network["client_bandwidth_limit"] < 0:
                print(
                    "{}Client bandwidth limit for network {} must be an integer greater than or equal to 0!{}".format(
                        Fore.RED, nname, Style.RESET_ALL
                    )
                )
                nerrors += 1
                continue

            shape_ret = net_obj.apply_shaping(int(network["client_bandwidth_limit"]))
            if shape_ret:
                print("{}update: added traffic shaping to network {}{}".format(Fore.YELLOW, nname, Style.RESET_ALL))
            else:
                print("{}Error adding shaping to network {}!{}".format(Fore.RED, nname, Style.RESET_ALL))
                nerrors += 1

        if "devices" in network:
            for serial, dev in network["devices"].items():
                if "name" not in dev:
                    print("{}Invalid device {}: name is missing!{}".format(Fore.RED, serial, Style.RESET_ALL))
                    nerrors += 1
                    continue

                inv_dev = [device for device in inv if device["serial"] == serial]
                if len(inv_dev) == 1:
                    dev_obj = None
                    if inv_dev[0]["networkId"] is not None and inv_dev[0]["networkId"] != net_obj.get("id"):
                        try:
                            inv_net_obj = Network(key=config["api_key"], id=inv_dev[0]["networkId"])
                            dev_obj = Device(key=config["api_key"], id=inv_dev[0]["serial"], net=inv_net_obj)

                            res = dev_obj.remove_device()
                            if not res:
                                print(
                                    "{}Error removing {} from network {}!{}".format(
                                        Fore.RED, inv_dev[0]["serial"], inv_dev[0]["networkId"], Style.RESET_ALL
                                    )
                                )
                                nerrors += 1
                                continue
                            print(
                                "{}update: removed {} from network {}{}".format(
                                    Fore.YELLOW, inv_dev[0]["serial"], inv_dev[0]["networkId"], Style.RESET_ALL
                                )
                            )
                            res = net_obj.claim_device(dev_obj)
                            if not res:
                                print("{}Error claiming {}!{}".format(Fore.RED, inv_dev[0]["serial"], Style.RESET_ALL))
                                nerrors += 1
                                continue
                            print("{}update: claimed {}{}".format(Fore.YELLOW, inv_dev[0]["serial"], Style.RESET_ALL))
                            dev_obj = res
                        except Exception as e:
                            print(
                                "{}Error updating device network membership for {}: {}{}".format(
                                    Fore.RED, inv_dev[0]["serial"], e, Style.RESET_ALL
                                )
                            )
                            nerrors += 1
                            continue
                    elif inv_dev[0]["networkId"] is None:
                        try:
                            dev_obj = Device(key=config["api_key"], id=inv_dev[0]["serial"], net=net_obj)
                            res = net_obj.claim_device(dev_obj)
                            if not res:
                                print("{}Error claiming device {}{}".format(Fore.RED, inv_dev[0]["serial"], Style.RESET_ALL))
                                nerrors += 1
                                continue
                            print("{}update: claimed {}{}".format(Fore.YELLOW, inv_dev[0]["serial"], Style.RESET_ALL))
                            dev_obj = res
                        except Exception as e:
                            print("{}Error claiming device {}: {}{}".format(Fore.RED, inv_dev[0]["serial"], e, Style.RESET_ALL))
                            nerrors += 1
                            continue

                    else:
                        dev_obj = Device(key=config["api_key"], id=inv_dev[0]["serial"], net=net_obj)
                        print("{}ok: {} is in network{}".format(Fore.GREEN, inv_dev[0]["serial"], Style.RESET_ALL))

                    dev_location = network["address"]
                    dev_name = dev["name"]
                    dparams = {"name": dev_name, "move_map_marker": True}
                    if "location" in dev:
                        dev_location += "\n" + dev["location"]
                        dparams["address"] = dev_location
                    if "tags" in dev:
                        dparams["tags"] = dev["tags"]
                    dev_obj.update_device(**dparams)
                    print("{}update: updated {} parameters{}".format(Fore.YELLOW, inv_dev[0]["serial"], Style.RESET_ALL))

                else:
                    print("{}Error finding {} in inventory!{}".format(Fore.RED, serial, Style.RESET_ALL))
                    nerrors += 1

        if "vlans" in network:
            vres = net_obj.enable_vlans()
            if not vres:
                print("{}Error enabling VLANs for network {}!{}".format(Fore.RED, nname, Style.RESET_ALL))
                nerrors += 1
                continue
            else:
                print("{}update: enabled VLANs for network {}{}".format(Fore.YELLOW, nname, Style.RESET_ALL))

            for vname, vlan in network["vlans"].items():
                done_msg = ""
                if int(vlan["id"]) != 1:
                    vlan_obj = net_obj.create_vlan(vname, vlan["id"], vlan["subnet"], vlan["appliance_ip"])
                    done_msg = "{}update: created VLAN {} (id={}, subnet={}, appliance_ip={}){}".format(
                        Fore.YELLOW, vname, vlan["id"], vlan["subnet"], vlan["appliance_ip"], Style.RESET_ALL
                    )
                else:
                    vlan_obj = Vlan(key=config["api_key"], id=1, net=net_obj)
                    done_msg = "{}ok: VLAN with ID {} exists{}".format(Fore.GREEN, vlan["id"], Style.RESET_ALL)
                if vlan_obj is None:
                    print(
                        "{}Error creating VLAN {} (id={}, subnet={}, appliance_ip={})!{}".format(
                            Fore.RED, vname, vlan["id"], vlan["subnet"], vlan["appliance_ip"], Style.RESET_ALL
                        )
                    )
                    nerrors += 1
                    continue
                print(done_msg)
                vargs = {}
                for key in ["reserved_ip_ranges", "fixed_ip_assignments", "dns_nameservers"]:
                    if key in vlan:
                        vargs[key] = vlan[key]
                res = vlan_obj.update_vlan(**vargs)
                vargs_str = ", ".join(["{}={}".format(k, v) for k, v in vargs.items()])
                if not res:
                    print("{}Error updating VLAN {} ({})!{}".format(Fore.RED, vname, vargs_str, Style.RESET_ALL))
                    nerrors += 1
                else:
                    print("{}update: Update VLAN {} ({}){}".format(Fore.YELLOW, vname, vargs_str, Style.RESET_ALL))

        if "ssids" in network:
            if len(network["ssids"].keys()) > 15:
                print("{}Only fifteen SSIDs are allowed per network!{}".format(Fore.RED, Style.RESET_ALL))
                nerrors += 1
            else:
                si = 0
                for sname, ssid in network["ssids"].items():
                    ssid_obj = SSID(key=config["api_key"], id=si, name=sname, net=net_obj)
                    sargs = {}
                    for key in ["name", "enabled", "auth_mode", "encryption_mode", "psk", "ip_assignment_mode", "tags"]:
                        if key in ssid:
                            sargs[key] = ssid[key]
                    res = ssid_obj.update_ssid(**sargs)
                    sargs_str = ", ".join(["{}={}".format(k, v) for k, v in sargs.items()])
                    if not res:
                        print("{}Error updating SSID {} ({})!{}".format(Fore.RED, sname, sargs_str, Style.RESET_ALL))
                        nerrors += 1
                    else:
                        print("{}update: Update SSID {} ({}){}".format(Fore.YELLOW, sname, sargs_str, Style.RESET_ALL))

                    if "allow_lan_access" in ssid and ssid["allow_lan_access"]:
                        sres = ssid_obj.allow_local_lan()
                        if not sres:
                            print("{}Error allowing local LAN access for SSID {}!{}".format(Fore.RED, sname, Style.RESET_ALL))
                            nerrors += 1
                        else:
                            print("{}update: Allowing local LAN access for SSID {}{}".format(Fore.YELLOW, sname, Style.RESET_ALL))

                    si += 1

        if "switches" in network:
            for serial, switch in network["switches"].items():
                dev_obj = Device(key=config["api_key"], id=serial, net=net_obj)
                if not dev_obj.realize():
                    print("{}Device {} is not in network {}{}".format(Fore.RED, serial, net_obj.get("name"), Style.RESET_ALL))
                    nerrors += 1
                    continue

                for port_range, switchport in switch.items():
                    ports = []
                    if isinstance(port_range, int) or port_range.isnumeric():
                        port_obj = SwitchPort(key=config["api_key"], id=port_range, dev=dev_obj)
                        ports.append(port_obj)
                    else:
                        prs = port_range.split(",")
                        for pr in prs:
                            pr = pr.strip()
                            if pr.isnumeric():
                                port_obj = SwitchPort(key=config["api_key"], id=pr, dev=dev_obj)
                                ports.append(pr)
                            else:
                                if "-" not in pr:
                                    print("{}Port range {} is invalid.{}".format(Fore.RED, pr, Style.RESET_ALL))
                                    nerrors += 1
                                    continue

                                (start, end) = pr.split("-")
                                start = start.strip()
                                end = end.strip()
                                if not start.isnumeric() or not end.isnumeric():
                                    print(
                                        "{}Error with port range '{}', {} and {} must be integers{}".format(
                                            Fore.RED, pr, start, end, Style.RESET_ALL
                                        )
                                    )
                                    nerrors += 1
                                    continue

                                if start >= end:
                                    print("{}Error with port range {}; start must be less than end{}".format(Fore.RED, pr, Style.RESET_ALL))
                                    nerrors += 1
                                    continue

                                pi = start
                                while pi <= end:
                                    port_obj = SwitchPort(key=config["api_key"], id=pi, dev=dev_obj)
                                    ports.append(port_obj)
                                    pi += 1

                    for port in ports:
                        pargs = {}
                        for key in [
                            "name",
                            "tags",
                            "enabled",
                            "type",
                            "vlan",
                            "voice_vlan",
                            "allowed_vlans",
                            "rstp_enabled",
                            "poe_enabled",
                        ]:
                            if key in switchport:
                                pargs[key] = switchport[key]
                        res = port.update_switchport(**pargs)
                        pargs_str = ", ".join(["{}={}".format(k, v) for k, v in pargs.items()])
                        if not res:
                            print("{}Error updating switchport range {} ({}){}".format(Fore.RED, port_range, pargs_str, Style.RESET_ALL))
                            nerrors += 1
                        else:
                            print("{}update: Update switchport range {} ({}){}".format(Fore.YELLOW, port_range, pargs_str, Style.RESET_ALL))

        if nerrors == 0:
            print("{}ok: network {} has been setup successfully!{}\n".format(Fore.GREEN, nname, Style.RESET_ALL))
        else:
            print(
                "{}Error fully configuring network {}.  See the errors above for more details.{}\n".format(Fore.RED, nname, Style.RESET_ALL)
            )
            errors += nerrors

    if errors == 0:
        print("{}ok: all networks have been setup successfully!{}\n".format(Fore.GREEN, Style.RESET_ALL))
    else:
        print(
            "{}There were errors setting up some of the networks.  See the output above for more details.{}\n".format(
                Fore.RED, Style.RESET_ALL
            )
        )
        exit(1)


if __name__ == "__main__":
    main()