#!/usr/bin/env python3

import argparse
import subprocess

def generate_sshfp_records(hostname, ip, local):
    '''Given a hostname and and IP address, scan the IP address (hostname
    not in dns yet) and return a bind string with sshfp records'''

    if local:
        p = ['ssh-keyscan', '-D', ip]
    else:
        # Handle being run via sudo which is the usual way
        # this is run.
        p = ['ssh', '-o', 'StrictHostKeyChecking=no',
             '-i', '/root/.ssh/id_rsa',
             'root@%s' % ip, 'ssh-keygen', '-r', ip]

    s = subprocess.run(p,
                       stdout=subprocess.PIPE,
                       stderr=subprocess.PIPE).stdout.decode('utf-8')
    fingerprints = []
    for line in s.split('\n'):
        if not line:
            continue
        _, _, _, algo, key_type, fingerprint = line.split(' ')
        # ssh-keygen on the host seems to return DSS/DSA keys, which
        # aren't valid to log in and not shown by ssh-keyscan -D
        # ... prune it.
        if algo == '2':
            continue
        fingerprints.append(
            (algo, key_type, fingerprint))

    # sort by algo and key_type to keep it consistent
    fingerprints = sorted(fingerprints,
                          key=lambda x: (x[0], x[1]))

    ret = ''
    first = True
    for f in fingerprints:
        ret += '%s%s\t\tIN\tSSHFP\t%s %s %s' % \
            ("\n" if not first else '', hostname, f[0], f[1], f[2])
        first = False
    return ret


def sshfp_print_records(hostname, ip, local=False):
    print(generate_sshfp_records(hostname, ip, local))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("hostname", help="hostname")
    parser.add_argument("ip", help="address to scan")
    parser.add_argument("--local", action='store_true',
                        help="Run keyscan locally, rather than via ssh")
    args = parser.parse_args()

    sshfp_print_records(args.hostname, args.ip, args.local)

if __name__ == '__main__':
    main()