Permalink
Browse files

Fix some errors. Rewrite all output. Refactoring.

  • Loading branch information...
1 parent 64742df commit 02e632d9f46b9e76d1c4a461820c815c59b9be37 @zzzsochi committed Dec 27, 2012
Showing with 121 additions and 43 deletions.
  1. +1 −1 setup.py
  2. +1 −1 ssh_authorizer/__main__.py
  3. +56 −13 ssh_authorizer/commands.py
  4. +63 −28 ssh_authorizer/helpers.py
View
@@ -2,7 +2,7 @@
setup(
name='ssh_authorizer',
- version='1.1',
+ version='1.2',
description='Manager for remote ~/.ssh/authorized_keys.',
author='Alexander Zelenyak aka ZZZ',
author_email='zzz.sochi@gmail.com',
@@ -13,7 +13,7 @@
from .helpers import parse_ssh_string
from .commands import help, get, add, delete, test
-logging.basicConfig(level=logging.ERROR)
+logging.basicConfig(format='%(message)s', level=logging.INFO)
commands = ['help', 'get', 'add', 'del', 'test']
@@ -1,9 +1,13 @@
import sys
+import logging
import sh
from .helpers import SSHController, SCPController
-from .helpers import get_authorized_keys, set_authorized_keys
+from .helpers import NoSuchFileError
+from .helpers import get_authorized_keys
+from .helpers import set_authorized_keys
+from .helpers import create_authorized_keys_file
from .helpers import load_local_keys
@@ -14,25 +18,35 @@ def help():
def get(user, host, port, raw):
if not raw:
+ ssh_controller = SSHController(user, host, port)
+
try:
- keys = [k for k in get_authorized_keys(user=user, host=host, port=port) if k]
+ keys = [k for k in get_authorized_keys(controller=ssh_controller) if k]
except sh.ErrorReturnCode_1:
sys.exit(1)
+ except NoSuchFileError:
+ keys = []
- if keys:
- print("Found {} keys:\n".format(len(keys)))
+ if len(keys) == 1:
+ print("{c.user}@{c.host}:{c.port} - found one key:\n".format(c=ssh_controller))
+ elif keys:
+ print("{c.user}@{c.host}:{c.port} - found {} keys:\n".format(len(keys), c=ssh_controller))
else:
- print("Not found keys.")
+ print("{c.user}@{c.host}:{c.port} - not found keys".format(c=ssh_controller))
for n, key in enumerate(keys):
if key:
print('{}: {}'.format(n + 1, key))
else:
+ ssh_controller = SSHController(user, host, port)
+
try:
- keys = get_authorized_keys(user=user, host=host, port=port)
+ keys = get_authorized_keys(controller=ssh_controller)
except sh.ErrorReturnCode_1:
sys.exit(1)
+ except NoSuchFileError:
+ keys = []
print('\n'.join(keys))
@@ -45,17 +59,27 @@ def add(user, host, port, key_files):
ssh_controller = SSHController(user, host, port)
try:
- remote_keys = get_authorized_keys(controller=ssh_controller)
+ remote_keys = get_authorized_keys(ssh_controller)
except sh.ErrorReturnCode_1:
sys.exit(1)
+ except NoSuchFileError:
+ create_authorized_keys_file(ssh_controller)
+ remote_keys = []
new_keys = []
+ already_keys = []
for key_file in key_files:
key = local_keys[key_file]
if key not in remote_keys:
new_keys.append(key)
+ else:
+ already_keys.append(key_file)
+
+ if already_keys:
+ logging.info('{c.user}@{c.host}:{c.port} - already in authorized_keys: "{}"'
+ .format('", "'.join(already_keys), c=ssh_controller))
if new_keys:
keys = remote_keys + new_keys
@@ -64,26 +88,45 @@ def add(user, host, port, key_files):
scp_controller.password = ssh_controller.password
try:
- set_authorized_keys(keys, controller=scp_controller)
+ set_authorized_keys(scp_controller, keys)
except sh.ErrorReturnCode_1:
sys.exit(1)
def delete(user, host, port, key_ids):
ssh_controller = SSHController(user, host, port)
- remote_keys = [k for k in get_authorized_keys(controller=ssh_controller) if k]
- for key_id in sorted(key_ids, reverse=True):
- del remote_keys[key_id - 1]
+ try:
+ remote_keys = [k for k in get_authorized_keys(ssh_controller) if k]
+ except NoSuchFileError:
+ logging.critical('{c.user}@{c.host}:{c.port} - error: not found authorized_keys'
+ .format(c=ssh_controller))
+ sys.exit(1)
+
+ try:
+ for key_id in sorted(key_ids, reverse=True):
+ del remote_keys[key_id - 1]
+ except IndexError:
+ logging.critical('{c.user}@{c.host}:{c.port} - error: not found key indexes'
+ .format(c=ssh_controller))
+ sys.exit(1)
scp_controller = SCPController(user, host, port)
scp_controller.password = ssh_controller.password
- set_authorized_keys(remote_keys, controller=scp_controller)
+ set_authorized_keys(scp_controller, remote_keys)
def test(user, host, port, key_files):
local_keys = load_local_keys(key_files)
- remote_keys = [k for k in get_authorized_keys(user=user, host=host, port=port) if k]
+
+ ssh_controller = SSHController(user, host, port)
+
+ try:
+ remote_keys = [k for k in get_authorized_keys(ssh_controller) if k]
+ except NoSuchFileError:
+ logging.info('{c.user}@{c.host}:{c.port} - not found authorized_keys'
+ .format(c=ssh_controller))
+ remote_keys = []
oks = []
@@ -3,7 +3,11 @@
from getpass import getpass
from tempfile import NamedTemporaryFile
-from sh import ssh, scp
+from sh import ssh, scp, ErrorReturnCode_1
+
+
+class NoSuchFileError(Exception):
+ pass
def parse_ssh_string(ssh_string):
@@ -40,9 +44,11 @@ def parse_ssh_string(ssh_string):
def load_local_keys(key_files):
- logging.debug('load local keys: {}'.format(key_files))
if not key_files:
key_files.append(os.path.expanduser('~/.ssh/id_rsa.pub'))
+ logging.info('Loading local id_rsa.pub')
+ else:
+ logging.info('Loading keys: {}'.format(', '.join(key_files)))
local_keys = {}
@@ -51,47 +57,41 @@ def load_local_keys(key_files):
key_data = f.read().strip()
local_keys[key_file] = key_data
- logging.debug('loaded local keys: {}'.format(local_keys))
return local_keys
class Controller(object):
out = b''
- error = b''
- password = None
user = None
host = None
port = 22
+ password = None
def __init__(self, user, host, port=22):
self.user = user
self.host = host
self.port = port or 22
def __call__(self, *args, **kwargs):
- logging.info('run command: "{}"'.format(self.process.ran))
+ logging.debug('run command: "{}"'.format(self.process.ran))
- def out_iteract(self, char, stdin):
+ def out_iteract(self, char, stdin, process):
if isinstance(char, str):
self.out += char.encode('utf8')
else:
self.out += char
- if self.out.decode('utf-8', errors='ignore').endswith('password: '):
+ out = self.out.decode('utf-8', errors='ignore')
+
+ if out.endswith('password: '):
self.clear()
stdin.put(self.get_password() + '\n')
- def err_iteract(self, char, stdin):
- if isinstance(char, str):
- self.out += char.encode('utf8')
- else:
- self.out += char
-
def get_password(self):
logging.debug('request password')
if not self.password:
- prompt = 'Need password for {}@{}:{}: '.format(self.user, self.host, self.port)
+ prompt = '{c.user}@{c.host}:{c.port} - need password: '.format(c=self)
self.password = getpass(prompt)
return self.password
@@ -105,18 +105,29 @@ def wait(self):
class SSHController(Controller):
+ no_such_file_error = False
+
def __call__(self, *args, **kwargs):
self.process = ssh(
'-o UserKnownHostsFile=/dev/null',
'-o StrictHostKeyChecking=no',
'-o LogLevel=quiet',
'{}@{}'.format(self.user, self.host), '-p', self.port,
- *args,
+ 'LANG=C', *args,
_out=self.out_iteract, _out_bufsize=0, _tty_in=True,
- _err=self.err_iteract, **kwargs)
+ **kwargs)
super().__call__(*args, **kwargs)
+ def out_iteract(self, char, stdin, process):
+ super().out_iteract(char, stdin, process)
+
+ out = self.out.decode('utf-8', errors='ignore')
+
+ if out.endswith('No such file or directory'):
+ self.no_such_file_error = True
+ process.kill()
+
class SCPController(Controller):
def __call__(self, local_file, remote_file, **kwargs):
@@ -128,20 +139,26 @@ def __call__(self, local_file, remote_file, **kwargs):
local_file,
'{}@{}:{}'.format(self.user, self.host, remote_file),
_out=self.out_iteract, _out_bufsize=0, _tty_in=True,
- _err=self.err_iteract, **kwargs)
+ **kwargs)
super().__call__(local_file, remote_file, **kwargs)
-def get_authorized_keys(controller=None, user=None, host=None, port=None):
- if not controller and (user and host):
- controller = SSHController(user, host, port)
- elif not controller:
- raise ValueError('Must set controller or user and host.')
+def get_authorized_keys(controller):
+ logging.info('{c.user}@{c.host}:{c.port} - getting authorized_keys'.format(c=controller))
try:
+ controller.clear()
controller('cat ~/.ssh/authorized_keys')
controller.wait()
+
+ except ErrorReturnCode_1:
+ if controller.no_such_file_error:
+ raise NoSuchFileError()
+ else:
+ logging.critical(controller.out.decode('utf8', errors='ignore'))
+ raise
+
except Exception:
logging.critical(controller.out.decode('utf8', errors='ignore'))
raise
@@ -150,11 +167,28 @@ def get_authorized_keys(controller=None, user=None, host=None, port=None):
return [line.strip() for line in out.split('\n')]
-def set_authorized_keys(keys, controller=None, user=None, host=None, port=None):
- if not controller and (user and host):
- controller = SCPController(user, host, port)
- elif not controller:
- raise ValueError('Must set controller or user and host.')
+def create_authorized_keys_file(controller):
+ logging.info('{c.user}@{c.host}:{c.port} - creating ~/.ssh'.format(c=controller))
+
+ try:
+ controller.clear()
+ controller('mkdir -p ~/.ssh')
+ controller.wait()
+
+ except ErrorReturnCode_1:
+ if controller.no_such_file_error:
+ raise NoSuchFileError()
+ else:
+ logging.critical(controller.out.decode('utf8', errors='ignore'))
+ raise
+
+ except Exception:
+ logging.critical(controller.out.decode('utf8', errors='ignore'))
+ raise
+
+
+def set_authorized_keys(controller, keys):
+ logging.info('{c.user}@{c.host}:{c.port} - writing authorized_keys'.format(c=controller))
with NamedTemporaryFile('w+b', buffering=0) as tmp:
data = '\n'.join(keys)
@@ -164,6 +198,7 @@ def set_authorized_keys(keys, controller=None, user=None, host=None, port=None):
tmp.write(b'\n')
try:
+ controller.clear()
controller(tmp.name, '~/.ssh/authorized_keys')
controller.wait()
except Exception:

0 comments on commit 02e632d

Please sign in to comment.