sshcon/sshcon.py at master · avi3tal/sshcon · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
#!/usr/bin/python
import sys, os
import re

DIR = os.path.dirname(os.path.abspath(__file__))
# DIR = '/home/avi/bin/'
sshargs = '-X -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no'

def set_color (string, color):
    colors = {"default": 0, "red": 31, "green": 32, "yellow": 33, "blue": 34,
              "magenta": 35, "cyan": 36, "white": 37, "black": 39}  # 33[%colors%m

    color_string = "\033[%dm\033[1m" % colors[color]
    return color_string + string + '\033[0m'

def collectConn():
    groups = {}
    gname = 'global'
    groups[gname] = []
    with open(os.path.join(DIR, '.servers'), 'r') as f:
        for line in f.readlines():
            if line.strip() and line.startswith('#'):
                continue
            match = re.match(r'^\[(\w+)\]$', line.strip())
            if match:
                gname = match.groups()[0]
                groups[gname] = []
                continue
            if line.strip() and not line.startswith('#'):
                line = line.replace(' ', '').split(',')
                conn = line[0]
                if not conn.split('@')[0]:
                    conn = os.getenv('USERNAME') + line[0]
                line = [conn, line[1].strip()]
                groups[gname].append(line)
    if not groups['global']:
        groups.pop('global')
    return groups

def colorConnList(serversDict):
    newServersList = []
    colors = ['red', 'white', 'cyan', 'magenta', 'green', 'yellow']
    colorIndex = 0
    for grp, serversList in serversDict.iteritems():
        if colorIndex >= len(colors):
            colorIndex = 0
        for serv in serversList:
            newServersList.append(serv)
            msg = '%s - %s' % (newServersList.index(serv) + 1, serv[1])
            print set_color(msg, colors[colorIndex])
        colorIndex += 1
    return newServersList

def getServerData(server, serversList):
    for serv in serversList:
        if server in serv:
            return serv
    return 'root:%s' % server

def buildCommand(data):
    cmd = 'ssh %s %s'
    cmd = cmd % (sshargs, data[0])
    return cmd

def usage(servers):
    for grp, lconn in servers.iteritems():
        for conn in lconn:
            print conn[1]
    sys.exit(0)


if __name__ == '__main__':
    serversDict = collectConn()


    dest_info = ''
    serversList = []

    if len(sys.argv) > 1:
        if '-m' in sys.argv:
            usage(serversDict)
        dest_info = getServerData(sys.argv[1], serversList)
    else:
        serversList = colorConnList(serversDict)
        number = input("Enter Server number:")
        if number <= 0:
            print 'Negetive or Zero numbers are not supported.'
            sys.exit(1)
        number -= 1
        dest_info = serversList[number]

    os.system(buildCommand(dest_info))