mssh

#!/usr/bin/python
# -*- coding:utf-8 -*-

'''
File: mssh
Author: Wayne Xue <xwsou@gmail.com>
Date: 2014-11-11
Description: ManySSH main script
'''

import os
import sys
import re
import argparse
import pexpect
import json
from msshdb import Database

prog_nama = 'ManySSH'
prog_desc = 'A ssh connection management system'
prog_ver = '1.1'
prog_ver_num = 3

db = Database(prog_ver_num)


class ManySSH(object):
    """docstring for ManySSH"""
    def __init__(self):
        parser = argparse.ArgumentParser(prog='mssh', description=prog_desc)
        parser.add_argument('-l', '--list', action='store_true',
                            help='list connections')
        parser.add_argument('-s', '--search', metavar='query',
                            help='search connections')
        parser.add_argument('-d', '--detail', action='store_true',
                            help='show detail information when listing')
        parser.add_argument('-f', '--ftp', action='store_true',
                            help='using sftp')
        parser.add_argument('--add', metavar='connection_string',
                            help='add connection')
        parser.add_argument('--identity', metavar='identity_path',
                            help='identity path when adding a connection')
        parser.add_argument('--password', metavar='password',
                            help='password when adding a connection')
        parser.add_argument('--alias', metavar='alias',
                            help='alias when adding a connection')
        parser.add_argument('--edit', nargs=3, metavar='param',
                            help='edit connection field by id')
        parser.add_argument('--delete', type=int, metavar='connection_id',
                            help='delete connection by id')
        parser.add_argument('--json', action='store_true',
                            help='output data with json')
        parser.add_argument('-v', '--version', action='version',
                            version=prog_nama + " v" + prog_ver + " buildin " + str(prog_ver_num),
                            help='display version info')
        parser.add_argument('param', nargs='?',
                            help='connection alias, param string or id')
        args = parser.parse_args()
        self.parser = parser
        self.args = args

    def run(self):
        """ManySSH run"""
        args = self.args
        try:
            if args.list is True:
                self.list_connection()
            elif args.search is not None:
                self.search_connection()
            elif args.add is not None:
                self.add_connection()
            elif args.delete is not None:
                self.delete_connection()
            elif args.edit is not None:
                self.edit_connection()
            elif args.param is not None:
                self.connect()
            else:
                self.parser.print_usage()
        except KeyboardInterrupt, e:
            self.exit('')

    def print_connections(self, rs, json_type=False):
        """print many connections"""
        args = self.args
        if not json_type:
            print ''
            print 'ID\tALIAS\t\tCONNECTION'
            for r in rs:
                print self.make_connection_string(r)
            print ''
        else:
            data = []
            for r in rs:
                d = {"id": r[0], "alias": r[6], "user": r[2], "host": r[1], "port": r[3]}
                if args.detail:
                    d["password"] = r[4]
                    d["identify_file"] = r[5]
                data.append(d)
            print json.dumps(data)

    def make_connection_string(self, r):
        """a kind look of connection data"""
        args = self.args
        detail = ""
        if args.detail:
            detail = '  '
            detail += (str(r[5]) if r[5] is not None and r[5] != ""
                                else str(r[4]))
        return ("*" + str(r[0]) + '\t' + r[6] + '\t\t' + r[2] + '@' + r[1] + ':'
                + str(r[3]) + detail)

    def connect(self):
        """connect the server"""
        args = self.args
        param = args.param
        if self.is_numberic(param):
            self.connect_by_id(param)
        else:
            self.connect_by_str(param)

    def connect_by_id(self, cid):
        """connect the server by connection id"""
        rs = db.find_connections(cid=cid)
        if len(rs) == 0:
            print 'connection(id: ' + str(cid) + ') not found'
        else:
            self.do_connect(rs[0])

    def connect_by_str(self, param):
        rs = db.find_connections(alias=param)
        if len(rs) > 0:
            self.do_connect(rs[0])
        else:
            user, host, port = self.get_param(param)
            '''select order:
                host
                host user
                host user port
            '''
            if user == '':
                port = 0
            rs = db.find_connections(host, user, port)
            if len(rs) > 1:
                print 'More than one connections found'
                self.print_connections(rs)
            elif len(rs) == 0:
                print 'No connection matched'
            else:
                self.do_connect(rs[0])

    def do_connect(self, r):
        """using connection data to connect"""
        args = self.args
        cstr = self.make_connection_string(r)
        print "Connect to " + cstr,
        cmd_args = []
        if not args.ftp:
            cmd = 'ssh'
            cmd_args += ['-o', 'ServerAliveInterval=30', '-p', str(r[3])]
        else:
            cmd = 'sftp'
            cmd_args += ['-P', str(r[3])]
        if r[5] != '':
            cmd_args += ['-i', r[5]]
        cmd_args += [r[2] + "@" + r[1]]
        child = pexpect.spawn(cmd, cmd_args)
        try:
            expects = ['password:', 'Connection refused', 'login from', 'Last login:', 'Are you sure you want to continue connecting (yes/no)?']
            index = child.expect(expects)
            if index == 0:
                child.sendline(r[4])
            elif index == 1:
                print ': Connection refused',
            elif index == 3:
                print '\nLast login:'
            elif index == 4:
                print '\nAre you sure you want to continue connecting (yes/no)?',
        except pexpect.EOF:
            print '\nEOF ERROR!'
        except pexpect.TIMEOUT:
            print '\nTIMEOUT!'
        child.interact()

    def list_connection(self):
        """list connections"""
        args = self.args
        kw = args.param
        if kw is None:
            kw = ''
        rs = db.find_connections(kw)
        if len(rs) > 0:
            self.print_connections(rs, args.json)
        else:
            print 'No connection found.'

    def search_connection(self):
        """list connections"""
        args = self.args
        q = args.search
        rs = db.search_connections(q)
        if len(rs) > 0:
            self.print_connections(rs, args.json)
        else:
            print 'No connection found.'

    def add_connection(self):
        """add connection"""
        args = self.args
        user, host, port = self.get_param(args.add)
        password = self.get_password()
        identity = self.get_identity()
        alias = self.get_alias()

        if len(db.find_connections(host, user, port)) > 0:
            self.exit('connection has been existed', 0)

        if user == '':
            self.exit('you must set a user for the connection')
        if identity == '' and password is None:
            self.exit('you must assign either identity file path or password')
        elif identity != '':
            if not os.path.exists(identity):
                self.exit('identity file does not exist')
            password = ''
        else:
            i = 0
            while password == '':
                print 'Connetion Password: '
                password = sys.stdin.readline().strip()
                i += 1
                if i > 2:
                    self.exit('Aboard', 0)
        if len(db.find_connections(host, user, port)) == 0:
            db.add_connection(host, user, password, port, identity, alias)
            print 'Added'
        else:
            self.exit('connection has been existed', 0)

    def delete_connection(self):
        """delete connection by id"""
        args = self.args
        rs = db.find_connections(cid=args.delete)
        if len(rs) == 0:
            print 'connection(id: ' + str(args.delete) + ') not found'
        else:
            db.delete_connection_by_id(args.delete)
            print 'Removed'

    def edit_connection(self):
        """edit connection by id"""
        args = self.args
        allows = ['host', 'user', 'port', 'password', 'idfile', 'alias']
        cid = args.edit[0]
        field = args.edit[1]
        value = args.edit[2]
        if not self.is_numberic(cid):
            self.exit('invalid connection id: ' + cid)
        if value is None or value == "":
            self.exit('new value can not be blank')
        rs = db.find_connections(cid=cid)
        if len(rs) == 0:
            self.exit('connection(id: ' + str(args.delete) + ') not found')
        r = db.get_connection_dict(rs[0])
        if not field in allows:
            self.exit('only one in the list ' + str(allows) + ' is allowed')
        if field in ['host', 'user', 'port']:
            r[field] = value
        rs = db.find_connections(r['host'], r['user'], r['port'], ecid=r['id'])
        if len(rs) > 0:
            print 'connection has been existed'
        else:
            db.edit_connection_by_id(cid, field, value)
            print 'Changed'

    def exit(self, msg, code=1):
        print msg
        sys.exit(code)

    def get_param(self, param):
        """split param string to user host and port"""
        # host & user
        if param.find('@') != -1:
            user, more = param.split('@', 1)
        else:
            user = ''
            more = param
        # host & port
        if more.find(':') != -1:
            host, port = more.split(':')
        else:
            host = more
            port = 22
        return user, host, port


    def get_password(self):
        """get password param"""
        args = self.args
        return args.password if args.password is not None else ''

    def get_identity(self):
        """get identity file path  param"""
        args = self.args
        return args.identity if args.identity is not None else ''

    def get_alias(self):
        """get alias param"""
        args = self.args
        return args.alias if args.alias is not None else ''

    def is_numberic(self, num):
        return re.compile(r'^\d+$').match(num) is not None


if __name__ == '__main__':
    ManySSH().run()