#!/usr/bin/env python3
# ssh-duphosts -- checks for duplicate entries in ~/.ssh/known_hosts
from operator import itemgetter
import argparse
import os
import sys

class Hostname(object):
    def __init__(self, value):
        self.value = value
        self.comparable = value.split(",", 1)[0]

    def __hash__(self):
        return hash(self.value)

    def __str__(self):
        return str(self.value)

    def __gt__(self, other):
        self_d = is_ip(self.comparable)
        other_d = is_ip(other.comparable)

        if self_d and not other_d:
            return True
        elif other_d and not self_d:
            return False
        else:
            return self.comparable > other.comparable

def is_ip(addr):
    if not addr:
        return False

    if addr[0] == "[":
        addr = addr[1:addr.find("]")]

    if ":" in addr:
        return True

    if all(x.isdigit() for x in addr.split(".")):
        return True

    return False

def find_duplicates(fh):
    keys = {}

    for line in fh:
        line = line.strip()
        if line == "" or line[0] == "#":
            continue

        try:
            if line[0] == "@":
                tag, host, ktype, key = line.split(" ", 3)
            else:
                host, ktype, key = line.split(" ", 2)
                tag = ""
        except ValueError:
            print("bad line %r" % line, file=sys.stderr)
            continue

        if (tag, ktype, key) in keys:
            keys[tag, ktype, key].append(host)
        else:
            keys[tag, ktype, key] = [host]

    return keys

def print_duplicates(keys):
    _keys = list(keys.keys())
    _keys.sort(key=itemgetter(2))
    _keys.sort(key=itemgetter(1))
    for entry in _keys:
        hosts = keys[entry]
        tag, ktype, key = entry
        if len(hosts) > 1:
            short_key = tag + (" " if tag else "") + ktype + " ..." + key[-15:]
            print("Key [%s] has %d entries:" % (short_key, len(hosts)))
            for host in hosts:
                addrs = host.split(",")
                print("\t%s" % "\n\t| ".join(addrs))

def print_merged(bykey, unmerge=False):
    byhost = {}

    for entry in bykey:
        tag, ktype, key = entry
        hosts = set()
        for item in bykey[entry]:
            hosts |= set(item.split(","))
        hosts.discard("")
        if unmerge:
            for host in hosts:
                host = Hostname(host)
                byhost[tag, host, ktype] = key
        else:
            hosts = sorted(hosts, key=Hostname)
            host = Hostname(",".join(hosts))
            byhost[tag, host, ktype] = key

    hosts = list(byhost.keys())
    hosts.sort(key=itemgetter(2))
    hosts.sort(key=itemgetter(1))

    for entry in hosts:
        tag, host, ktype = entry
        key = byhost[entry]
        if tag:
            print(tag, host, ktype, key)
        else:
            print(host, ktype, key)

opt_input = os.path.expanduser("~/.ssh/known_hosts")

parser = argparse.ArgumentParser()
parser.add_argument("-m", "--merge", action="store_true",
                    help="filter mode - merge entries with identical keys")
parser.add_argument("-M", "--unmerge", action="store_true",
                    help="filter mode - split one hostname per line")
parser.add_argument("input_file", nargs="?", default=opt_input,
                    help="known_hosts file to read")
args = parser.parse_args()

if args.merge + args.unmerge > 1:
    exit("error: merge and unmerge filters are mutually exclusive")

if args.input_file == "-":
    keys = find_duplicates(sys.stdin)
else:
    with open(args.input_file, "r") as fh:
        keys = find_duplicates(fh)

if args.merge or args.unmerge:
    print_merged(keys, args.unmerge)
else:
    print_duplicates(keys)
