~curtin-dev/curtin/xenial

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
#!/usr/bin/python3
import argparse
import subprocess
import os.path
import sys
import json
VALID_MODES = (
    # one key per line
    'dump',
    # suitable for use in cloud-config
    'cloud-config',
    # suitable for use in cloud-localds meta-data
    'meta-data',
    # simple json list of strings keys
    'json'
)


def jdump(data):
    return json.dumps(data, indent=2, sort_keys=True, separators=(',', ': '))


def read_pubkeys():
    keys = []
    try:
        out = subprocess.check_output(['ssh-add', '-L'],
                                      stderr=subprocess.STDOUT)
        if isinstance(out, bytes):
            out = out.decode()
        keys += out.splitlines()
    except Exception:
        pass

    ssh_dir = os.path.expanduser('~' + os.path.sep + '.ssh')
    keyfiles = ['id_ecdsa.pub', 'id_ed25519.pub', 'id_rsa.pub']
    for f in [os.path.join(ssh_dir, f) for f in keyfiles]:
        if os.path.isfile(f):
            with open(f, "r") as fp:
                keys += fp.read().splitlines()

    return list(set(keys))


def format_pubkeys(mode, keys):
    if mode not in VALID_MODES:
        raise ValueError("unknown mode '%s'. Expected one of: %s" %
                         (mode, ' '.join(VALID_MODES)))
    if mode == 'dump':
        if len(keys) != 0:
            data = '\n'.join(keys) + "\n"
    elif mode == "cloud-config" or mode == "meta-data":
        data = ""
        if mode == "cloud-config":
            data = "#cloud-config\n"
            name = 'ssh_authorized_keys'
        else:
            name = 'public-keys'
        data += name + ": " + jdump(keys) + "\n"
    elif mode == "json":
        return jdump(keys) + "\n"
    return data


def main():
    parser = argparse.ArgumentParser(
        description='Dump users public ssh keys in a variety of formats')
    parser.add_argument('mode', nargs='?', help='output mode',
                        choices=VALID_MODES, default='dump')
    args = parser.parse_args()
    sys.stdout.write(format_pubkeys(args.mode, read_pubkeys()))


if __name__ == '__main__':
    sys.exit(main())