0_setup.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import subprocess

import utils


def add_allowed_keys(path, keys):
    content = ''
    if os.path.exists(path):
        with open(path, 'r') as fo:
            content = fo.read()
    new_content = content.strip()
    for key in keys:
        if key not in new_content:
            new_content += '\n' + key
            utils.log('The key "%s" will be added in "%s".' % (key.split(' ')[-1], path))
        else:
            utils.log('The key "%s" is already in "%s".' % (key.split(' ')[-1], path))
    if new_content != content:
        with open(path, 'w') as fo:
            fo.write(new_content)
        utils.log('The file "%s" has been updated.' % path)
    else:
        utils.log('The file "%s" is already up to date.' % path)


def setup(interactive=True):
    dir_path = utils.get_dir(__file__)
    cmds = list()
    # Create / update ubicast account
    cmds.append('echo "Checking ubicast account"')
    code, out = utils.exec_cmd(['id', 'ubicast'])
    if code != 0:
        cmds.append('useradd -m -s /bin/bash ubicast')
        out = ''
    if 'sudo' not in out:
        cmds.append('usermod -aG sudo ubicast')
    # Add SSH key
    cmds.append('echo "Checking ubicast and root SSH keys"')
    allowed_keys = utils.get_conf('SSH_ALLOWED_KEYS', '').strip().split('\n')
    with open('%s/ubicast_support.pub' % dir_path, 'r') as fo:
        support_key = fo.read()
    allowed_keys.append(support_key.strip())
    # root
    cmds.append('mkdir -p /root/.ssh')
    cmds.append('chmod 700 /root/.ssh')
    add_allowed_keys('/root/.ssh/authorized_keys', allowed_keys)
    # ubicast
    cmds.append('mkdir -p /home/ubicast')
    cmds.append('mkdir -p /home/ubicast/.ssh')
    cmds.append('chmod 700 /home/ubicast/.ssh')
    add_allowed_keys('/home/ubicast/.ssh/authorized_keys', allowed_keys)
    cmds.append('cp "/root/.bashrc" "/home/ubicast/.bashrc"')
    cmds.append('chown -R ubicast:ubicast /home/ubicast/.ssh')

    utils.run_commands(cmds)
    # Set ubicast password if any
    pwd = utils.get_conf('SHELL_UBICAST_PWD')
    if pwd:
        p = subprocess.Popen(['passwd', '-q', 'ubicast'], stdin=subprocess.PIPE)
        p.communicate(input=b'%(pwd)s\n%(pwd)s' % {b'pwd': pwd.encode('utf-8')})
        if p.returncode != 0:
            raise Exception('Failed to set ubicast account password.')
        utils.log('\033[1;33m The ubicast account password has been set. \033[0m')