#!/usr/bin/python3
# Author: Alon Swartz <alon@turnkeylinux.org>

import os
import sys

if '_TURNKEY_INIT' in os.environ:
    sys.exit(0)

import pwd
import ec2metadata

USERNAME = 'root'


def authorize_sshkeys(keys, username):
    pw = pwd.getpwnam(username)
    sshdir = os.path.join(pw.pw_dir, '.ssh')
    authorized_keys = os.path.join(sshdir, 'authorized_keys')

    if not os.path.exists(sshdir):
        os.makedirs(sshdir)
        os.chmod(sshdir, 0o700)
        os.chown(sshdir, pw.pw_uid, pw.pw_gid)

    with open(authorized_keys, 'a') as fob:
        fob.write(''.join(['%s\n' % key for key in keys]))
    os.chown(authorized_keys, pw.pw_uid, pw.pw_gid)


def main():
    keys = ec2metadata.get('public-keys')
    if keys:
        authorize_sshkeys(keys, USERNAME)


if __name__ == "__main__":
    main()
