#!/usr/bin/env python
# vim:ts=2:et:sw=2:ai
#
# Find shortest path for SSH multiple access in case of broken routing 
#
# .bashrc snippet
# function _listnodes {
#         COMPREPLY=()
#         cur="${COMP_WORDS[COMP_CWORD]}"
# 	opts=$(find $HOME/wleiden/genesis/nodes/ -mindepth 1 -type d | xargs -n 1 basename)
#         COMPREPLY=( $(compgen -W "${opts}" -- ${cur}) )
#         return 0
# }
# complete -F _listnodes find_path.py
# export PATH=$PATH:$HOME/wleiden/genesis/tools
#
#
# Rick van der Zwet <info@rickvanderzwet.nl>
#
import argparse
import gformat
import glob
import os
import re
import subprocess
import sys
import time
import tempfile
import yaml

from collections import defaultdict

TMPFILE = '/tmp/test.dot'
OUTFILE = os.path.join(os.getcwd(),'network.png')
LINK_SPECIFIC = False

gateways = defaultdict(list)
neighbors = defaultdict(list)
datadumps = {}
nodes = {}
nodes_down = []
poel = {}



t_start = time.time()
def oprint(*args, **kwargs):
  print(">>> [%05.2f] " % (time.time() - t_start), end='', file=sys.stderr)
  print(*args, *kwargs, file=sys.stderr)



def generate_network_tree():
  link_type = {}
  link_data = {}

  oprint("# Load configuration files")
  try:
    for host in gformat.get_hostlist():
      datadump = gformat.get_yaml(host)
      datadumps[host] = datadump

      if not datadump['status'] == 'up':
        nodes_down.append(host)
        if not args.dry_run:
          continue

      nodename = datadump['nodename']
      iface_keys = [elem for elem in list(datadump.keys()) if (elem.startswith('iface_') and not "lo0" in elem)]
      nodes[nodename] = datadump
      for iface_key in iface_keys:
        # Bridge linked interfaces do not have IP address by themself
        if not 'ip' in datadump[iface_key]:
          continue

        # Annotation of device, no real key
        if '_extra' in iface_key:
          continue

        if not datadump[iface_key]['status'] == 'up':
          if not args.dry_run:
            continue

        l = datadump[iface_key]['ip']
        addr, mask = l.split('/')

        # Not parsing of these folks please
        if not gformat.valid_addr(addr):
          continue

        addr = gformat.parseaddr(addr)
        mask = int(mask)
        addr = addr & ~((1 << (32 - mask)) - 1)
        nk = (nodename, iface_key)
        if addr in poel: 
          poel[addr] += [nk]
        else: 
          poel[addr] = [nk]
          # Assume all ns_ip to be 11a for a moment
          if 'ns_ip' in datadump[iface_key]:
            link_type[addr] = '11a'
          else:
            link_type[addr] = datadump[iface_key]['type']

          link_data[addr] = 1
          iface = datadump[iface_key]['autogen_ifname']
          #print("### ... %s %s [%s] is of type %s" % (iface, gformat.showaddr(addr), iface_key, link_type[addr]))
  except (KeyError, ValueError) as e:
    print("[FOUT] in '%s' interface '%s'" % (host,iface_key)) 
    print(e)
    sys.exit(1)

  oprint("# Prepare buren dictionary") 
  for addr,leden in poel.items():
    leden = sorted(set(leden))
    for index,lid in enumerate(leden[:-1]):
      for buur in leden[index + 1:]:
        neighbors[lid[0]].append(buur[0])
        neighbors[buur[0]].append(lid[0])


def find_shortest_paths(src):
  hop_cnt = 0
  shortest_paths = defaultdict(list)
  shortest_paths[src] = [src]
  new_paths = [[src]]

  # List current path
  # Expand paths with new neighbors
  while new_paths:
    hop_cnt += 1
    path_canidates = new_paths
    new_paths = []

    for path_canidate in path_canidates:
      for neighbor in sorted(neighbors[path_canidate[-1]]):
        new_path = path_canidate + [neighbor]
        # Prune paths which are longer than already known path
        if neighbor in shortest_paths and len(shortest_paths[neighbor]) < len(new_path):
          continue

        # Register new shortest path  
        shortest_paths[neighbor] = new_path

        # Claim path for expansion
        new_paths.append(new_path)

  return(shortest_paths)



def find_shortest_path(src, dst):
  shortest_paths = find_shortest_paths(src)
  return shortest_paths[dst]



def show_path(path):
  oprint("#", " -> ".join(path))

  stack = path.copy()
  src = stack.pop(0)
  ip_path = [datadumps[src]['masterip']]
  # Walk path to find IP address of links
  while stack:
    dst = stack[0]
    for neighbors in poel.values():
      nodes = ([x[0] for x in neighbors])
      if src in nodes and dst in nodes:
        dst_iface_key = dict(neighbors)[dst]
        dst_iface_ip = datadumps[dst][dst_iface_key]['ip'].split('/')[0]
        ip_path.append(dst_iface_ip)
        break
    src = stack.pop(0)

  for i,(host,ip) in enumerate(zip(path, ip_path)):
    oprint("# %i: ... %s (%s)" % (i, host, ip))
    
  oprint("")

  if args.external:
    (first_hop, ext_port) = datadumps[path[0]]["remote_access"].split(':')
    ip_path.pop(0)
    ip_path.insert(0, first_hop)
  else:
    ext_port = 22

  cmd = "ssh -t -L 0.0.0.0:%(port)s:127.0.0.1:%(port)s -p %(ext_port)s " + \
    " ssh -t -L %(port)s:127.0.0.1:%(port)s ".join(ip_path) + \
    " ssh -t -D %(port)s 127.0.0.1"
  print(cmd % { 'port': args.proxy_port, 'ext_port': ext_port})
    
    

if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument('--dry-run', action='store_true', help="Include nodes and links which are down")
  parser.add_argument('--proxy-port', default=6256, help="Use unique portnumber to avoid clashes on multi-logins")
  parser.add_argument('--external', action='store_true', help="Access source node via internet remote access")
  parser.add_argument('src', help="Source node")
  parser.add_argument('dst', help="Destination node")
  args = parser.parse_args()

  
  generate_network_tree()
  oprint("# Find shortest path")
  path = find_shortest_path(args.src, args.dst)
  oprint("# Display results")
  show_path(path)
  oprint("# All done")