#!python
#
#   Copyright 2019 Jeremy Schulman, nwkautomaniac@gmail.com
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
"""
This script is used to locate an end Host that is attached on a layer 2 (VLAN)
network.  Prior to use the User must export the following environment variables:

   * PYATS_USERNAME=<network-device login user-name>
   * PYATS_PASSWORD=<login user-password>

The User is expected to provide the host IP address and the name of a Cisco
device to start the search on.  For example:

$ find-host -i 192.168.100.171 -o switch22 -F

INFO:Starting search for host IP 192.168.100.171 ...
INFO:Connecting to device: switch22
INFO:Ping host IP: 192.168.100.171
INFO:Found MAC address 0050.c2ab.abab on VLAN 99
INFO:Following switch22 interface Eth1/1 to next device: switch21
INFO:Host IP 192.168.100.171 macaddr 0050.c2ab.abab last found on switch39:Gi1/1/19
INFO:No next device available.
INFO:Building configuration...

Current configuration : 170 bytes
!
interface GigabitEthernet1/1/19
 description I-am-your-phone
 switchport access vlan 99
 switchport mode access
 spanning-tree portfast
 spanning-tree bpduguard enable
end
"""
import sys
from ipaddress import ip_address
import logging
import re
import textwrap

from first import first
import click

# unicon part of pyATS software
from unicon.core.errors import ConnectionError

import idreamofgenie

log = logging.getLogger(__name__)


def setup_log():
    """
    Using a log for processing/progress.  You can chance the format and controls
    in this function if you want to customize the logging output.
    """
    log.setLevel(logging.INFO)
    # formatter = logging.Formatter('%(asctime)s:%(levelname)s:%(message)s')
    formatter = logging.Formatter('%(levelname)s:%(message)s')
    ch = logging.StreamHandler(sys.stdout)
    ch.setFormatter(formatter)
    log.addHandler(ch)


setup_log()


# -----------------------------------------------------------------------------
#
#                       Click - User Interface Helpers
#
# -----------------------------------------------------------------------------

def validate_macaddr(ctx, param, value):
    """
    This is a Click callback that will validate the MAC address provided by the
    User.  No matter what MAC address format was provided by user, extract out
    only the valid alpha-num octets into a single string.  Ensure that the
    result is the proper length.

    Parameters
    ----------
    ctx : click.Context
    param : click.Option - the parameter (macaddr)
    value : str - the MAC address provided by the User

    Returns
    -------
    str
        The MAC address in Cisco format : "aaaa.bbbb.ccccc"
    """
    if not value:
        return value

    nibs = re.sub(r'[^0-9a-fA-F]', '', value.lower())

    if len(nibs) != 12:
        raise click.BadParameter('macaddr must be 6 octets in length')

    # return the macaddr value using the Cisco format of "aaaa.bbbb.cccc"
    return '.'.join(textwrap.wrap(nibs, 4))


def validate_ipaddr(ctx, param, value):
    """
    This is a Click callback that will validate the IP address provided by the
    User.

    Parameters
    ----------
    ctx : click.Context
    param : click.Option - the parameter (ipaddr)
    value : str - the IP address provided by the User

    Returns
    -------
    str
        The IP address
    """
    if not value:
        return value

    try:
        ip_address(value)
    except ValueError:
        click.BadParameter(f'`{value}` is not a valid IP address')

    return value


# -----------------------------------------------------------------------------
#
#                       Click - Command Line Interface
#
# -----------------------------------------------------------------------------

@click.command()
@click.option('-o', '--origin-device', required=True,
              help='Start find from this Cisco device')
@click.option('-n', '--os-name', type=click.Choice(['ios', 'iosxe', 'iosxr', 'nxos']),
              help='Origin device NOS name (if known)')
@click.option('-i', '--ipaddr',
              callback=validate_ipaddr,
              help='Target host IP address')
@click.option('-m', '--macaddr',
              callback=validate_macaddr,
              help='Target host MAC address')
@click.option('-F', '--follow', is_flag=True,
              help='Always follow next device, no prompt')
def find_host(ipaddr, macaddr, origin_device, os_name, follow):
    """
    This function is called via Click as the main CLI function.  The primary
    purpose is to ensure the User options are valid and then execute the
    code to find the host, traversing network devices starting at the 'origin switch'.

    Parameters
    ----------
    ipaddr : str - target Host IP address
    macaddr : str - target Host MAC address
    origin_device : str - origin Cisco switch
    os_name : str - origin switch network OS name
    follow : bool - auto-follow links
    """

    if not any((ipaddr, macaddr)):
        log.error("Required at least one of: [--ipaddr, --macaddr]")
        sys.exit(1)

    # check if the origin device was passed as an IP address rather than a
    # DNS/hostname; if it is not an IP address, then this function will attempt
    # to use the User provided value as a DNS/hostname.

    try:
        origin_device = ip_address(origin_device)
    except ValueError:
        pass

    try:
        log.info(f"Connecting to device: {origin_device}")
        device = idreamofgenie.Device(origin_device, os_name=os_name)

    except ConnectionError:
        log.error(f"Failed to connect to network device: {origin_device}")
        sys.exit(1)

    if macaddr:
        log.info(f"Searching for host with MAC address: {macaddr}")

    # if the target Host IP address was provided, execute a ping test to ensure
    # (a) that it is reachable, and (b) so that the target Host MAC address
    # will be available in the network device ARP tables.

    if ipaddr:
        log.info(f"Searching for host with IP address: {ipaddr}")
        log.info(f"Ping host IP: {ipaddr}")
        if not idreamofgenie.ping(device, ipaddr):
            log.error(f"Host IP {ipaddr} unreachable, exit.")
            sys.exit(1)

    # Start the hunt for the target device!  This function does not return
    # here; rather will execute the search and then exit the program using
    # sys.exit().

    hunt_end_host(device, ipaddr, macaddr, follow)


# -----------------------------------------------------------------------------
#
#                      Finding the Target Host Device
#
# -----------------------------------------------------------------------------

def hunt_end_host(device, ipaddr, macaddr, follow, vlan_id=None, ttl=4):
    """
    This function will examine the current `device` to see if the target can be
    found on a local interface.  If it is not on a local interface, then this
    function will connect to the next device where the Host was last seen (via
    ARP / MAC table / CDP), and then recursively call this function to continue
    the search.

    This function does not return, rather it will call sys.exit().

    Notes
    -----
    This function provides a Time-To-Live (ttl) counter that defaults to 4.
    This means if more than four devices are searched, and the Host is still
    not found, then this program will exit with an error.

    Parameters
    ----------
    device : Device - current device instance
    ipaddr : str|None - Host IP address
    macaddr : str|None - Host MAC address
    follow : bool - auto follow CDP links
    vlan_id : str|None - Host MAC found on VLAN ID
    ttl : int - time to live value
    """

    if not ttl:
        log.error(f"Count not find Host IP {ipaddr} in TTL, exit.")
        sys.exit(2)

    # If the Host MAC address is not known, then try to find it using an ARP
    # lookup using the Host IP address.  If the MAC address is not found, then
    # sys.exit with error.  If the MAC address is found, and the associated
    # interface is a VLAN, then set the vlan_id value which we use for
    # reporting purposes.

    if not macaddr:
        found = idreamofgenie.find_macaddr_by_arp(device, ipaddr)
        if not found:
            log.error(f"Could not find host MAC address from IP, exit.")
            sys.exit(2)
        macaddr = found['macaddr']
        found_vlan = found['interface'].lower().split('vlan')
        if len(found_vlan) > 1:
            vlan_id = found_vlan[1]

        if vlan_id:
            log.info(f"Found MAC address {macaddr} on VLAN {vlan_id}")

    # If the Host IP address is not known (User called with just --macaddr),
    # then try to find the IP address using the provided MAC address.  If it is
    # not found, we do not exit/error; but if it is found, then we save this
    # for further use.

    if not ipaddr:
        found = idreamofgenie.find_ipaddr_by_arp(device, macaddr)
        if found:
            log.info(f"Found IP address {ipaddr} for MAC address {macaddr}")
            ipaddr = found['ipaddr']

    # using the Host MAC address locate the device interface where it was last
    # seen using the mac address table.  From this via interface, gather the
    # current interface configuration as it will be presented to the User.

    via_iface = idreamofgenie.find_macaddr_via_iface(device, macaddr)
    iface_config = device.execute(f'show run int {via_iface}')

    # if the auto follow feature is disabled, show the User the current via
    # interface configuration and ask them if they want to continue searching.
    # if the auto follow feature is enabled, skip the User prompt.

    if follow is False:
        keep_hunting = click.confirm(f"Found MAC address {macaddr} on interface {via_iface}\n"
                                     f"{iface_config}\n\n"
                                     "Continue hunting?")
    else:
        keep_hunting = True

    if not keep_hunting:
        log.info(f"Host IP {ipaddr} macaddr {macaddr} last found on {via_iface}")
        log.info(iface_config)
        sys.exit(0)

    # if the via interface is a port-channel, then we need to obtain an actual
    # physical interface since we need to follow the links to the next device
    # in the hunt.

    if via_iface.lower().startswith('po'):
        members = idreamofgenie.find_portchan_members(device, via_iface)
        follow_iface = first(members)
        if not follow_iface:
            log.error(f"Could not follow Port-channel {via_iface}, no physical ports found")
            sys.exit(2)
    else:
        follow_iface = via_iface

    # Using the via interface and CDP, find the next Cisco device to search on.
    # If the next device does not have an os_name value, it means the next
    # device is not a Cisco device.

    next_device = idreamofgenie.find_cdp_neighbor(device, follow_iface)
    if not next_device or not next_device['os_name']:
        log.info(f"Host IP {ipaddr} macaddr {macaddr} last found on {device.name}:{via_iface}")
        log.info(f"No next device available.")
        log.info(iface_config)
        sys.exit(0)

    # At this point there is valid next Cisco device to search, update the User
    # and attempt to connect to that device.  Once connected, recursively call
    # this (hunt_end_host) function.

    next_name = next_device['device']
    log.info(f"Following {device.name} interface {follow_iface} to next device: {next_name}")

    try:
        device = idreamofgenie.Device(next_name, os_name=next_device['os_name'])

    except ConnectionError:
        log.error(f"Failed to connect to network device: {next_name}")
        sys.exit(1)

    hunt_end_host(device, ipaddr=ipaddr, macaddr=macaddr,
                  vlan_id=vlan_id,
                  follow=follow, ttl=ttl-1)


# -----------------------------------------------------------------------------
#
#                                   MAIN
#
# -----------------------------------------------------------------------------

if __name__ == "__main__":
    try:
        find_host()

    except Exception as exc:
        log.error(f"Failed to find target, error: {str(exc)}")
        sys.exit(1)
