django-fab-deploy-ext / fab_deploy_ext / ssh.py

import os
import re
import datetime

from fabric.api import get
from fabric.contrib.files import comment
from fabric.contrib.files import uncomment
from fabric.contrib.files import exists

from fab_deploy_ext import file
from fab_deploy_ext import utils


__all__ = ['ssh_list_files', 'ssh_list_keys', 'ssh_disable_key',
           'ssh_enable_key']


def _get_users(users, exclude_users=None):
    if exclude_users is None:
        exclude_users = []
    if users is None:
        if 'root' not in exclude_users:
            yield 'root'
        for dir in file.file_list('/home'):
            user = os.path.basename(dir)
            if user not in exclude_users:
                yield user
    else:
        if exclude_users:
            users = set(users) - set(exclude_users)
        for user in users:
            yield user


def _get_user_dir(user):
    if user == 'root':
        return '/root'
    else:
        return '/home/%s' % user


def ssh_list_files(users, exclude_users=None):
    users = _get_users(users, exclude_users=exclude_users)
    for user in users:
        dir = _get_user_dir(user)
        authorized_keys = '%s/.ssh/authorized_keys' % dir
        if exists(authorized_keys, use_sudo=True):
            yield user, authorized_keys


def ssh_list_keys(users=None, exclude_users=None,  include_comments=False,
                  use_sudo=False):
    for user, authorized_keys in ssh_list_files(users, exclude_users=exclude_users):
        content = file.file_read(authorized_keys, use_sudo=True)
        for ssh_key in utils.split_lines(content):
            if not include_comments and ssh_key.startswith('#'):
                continue
            yield user, ssh_key


def ssh_disable_key(key, users=None, exclude_users=None, use_sudo=False):
    regex = '%s' % re.escape(key)
    regex = regex.replace('\/', '/')

    now = datetime.datetime.now().strftime('%Y.%m.%d-%H.%M')
    backup = '.%s.bak' % now
    for user, authorized_keys in ssh_list_files(users, exclude_users=exclude_users):
        comment(authorized_keys, regex, use_sudo=use_sudo, backup=backup)


def ssh_enable_key(key, users=None, exclude_users=None, use_sudo=False):
    regex = '%s' % re.escape(key)

    now = datetime.datetime.now().strftime('%Y.%m.%d-%H.%M')
    backup = '.%s.bak' % now
    for user, authorized_keys in ssh_list_files(users, exclude_users=exclude_users):
        uncomment(authorized_keys, regex, use_sudo=use_sudo, backup=backup)
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.