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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
|
#!/usr/bin/python
# -*- coding: utf-8 -*-
#!/usr/bin/env python
# Copyright (C) 2013 Raúl Benencia <rul@kalgan.cc>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
# Little script to load SSH keys on demand. You can use it by writing a
# function like the following one in your shell rc file:
#
# odsa_ssh() {
# odsa "$*"
# /usr/bin/ssh "$*"
# }
#
# And then an alias to ssh
# alias ssh=odsa_ssh
#
# Blog post about this script: https://kalgan.cc/blog/posts/On_demand_ssh-add/
import signal
import sys
from os.path import expanduser
from paramiko import SSHConfig, Agent
from subprocess import call
def signal_handler(signal, frame):
print 'Caught ^C'
sys.exit(4)
class ODSA(): # This cryptic name means On Demand Ssh-Add
def __init__(self, cfg_filename=expanduser('~/.ssh/config')):
self.cfg_filename = cfg_filename
def keys_for_host(self, host):
"""
Return a list of paths to the private keys configured by the user for
logging in on the received host.
@return: a list of keys configured up in ~/.ssh/config
@rtype: list of paths to private keys
"""
try:
cfg_file = file(self.cfg_filename)
except IOError:
print "Couldn't open SSH config file: " + self.cfg_filename
sys.exit(3)
else:
config = SSHConfig()
config.parse(cfg_file)
keys = config.lookup(host).get('identityfile', None)
cfg_file.close()
if keys:
return [expanduser(k) for k in keys]
else:
return []
def get_public_key_from_file(self, key_filename):
"""
Return the public key in base64 format.
@return: a string containing a public key encoded in base64
@rtype: str
"""
try:
key_file = file(key_filename + ".pub")
except IOError:
print "Couldn't open the public key file " + key_filename + ".pub"
sys.exit(2)
else:
contents = key_file.read()
key_file.close()
return contents.split()[1]
def is_key_loaded(self, key_filename):
"""
Return whether a certain SSH key is loaded in the SSH agent
@return: True if the key is loaded, False otherwise.
@rtype: Boolean
"""
key = self.get_public_key_from_file(key_filename)
return any(key == k.get_base64() for k in Agent().get_keys())
def load_keys(self, keys):
"""
Spawn ssh_add command for loading the missing keys.
@return: Returns values of ssh_add command
@rtype: Int
"""
keys.insert(0, "/usr/bin/ssh-add")
return call(keys)
if __name__ == "__main__":
if len(sys.argv) == 1:
print 'Usage: ' + sys.argv[0] + ' HOST'
sys.exit(1)
signal.signal(signal.SIGINT, signal_handler)
odsa = ODSA()
enabled_keys = odsa.keys_for_host(sys.argv[1])
keys_to_load = filter(lambda k: not odsa.is_key_loaded(k), enabled_keys)
if keys_to_load:
odsa.load_keys(keys_to_load)
sys.exit(0)
|
