Skip to content
Browse files

support explicit --ssh-identity option

  • Loading branch information...
1 parent cdb3b7d commit 971b2824378b33989e81abd40ca960aad43d85e8 @lirazsiri lirazsiri committed Dec 20, 2012
Showing with 62 additions and 39 deletions.
  1. +2 −2 cloudtask/executor.py
  2. +43 −13 cloudtask/ssh.py
  3. +14 −23 cloudtask/task.py
  4. +3 −1 cloudtask/taskconf.py
View
4 cloudtask/executor.py
@@ -146,7 +146,7 @@ def callback():
raise self.Error(e)
try:
- self.ssh.copy_id(self.sshkey.public)
+ self.ssh.copy_id(self.sshkey)
if taskconf.overlay:
self.ssh.apply_overlay(taskconf.overlay)
@@ -167,7 +167,7 @@ def _cleanup(self):
if self.cleanup_command:
self.ssh.command(self.cleanup_command).close()
- self.ssh.remove_id(self.sshkey.public)
+ self.ssh.remove_id(self.sshkey)
except:
pass
View
56 cloudtask/ssh.py
@@ -9,9 +9,44 @@
# option) any later version.
#
+import os
from os.path import *
from command import Command
+from temp import TempFile
+import executil
+import hashlib
+
+class PrivateKey:
+ class Error(Exception):
+ pass
+
+ def __init__(self, path):
+ self.path = abspath(path)
+ if not isfile(self.path):
+ raise self.Error("no such file '%s'" % self.path)
+
+ @property
+ def public(self):
+ try:
+ return executil.getoutput("ssh-keygen -y -P '' -f", self.path)
+ except executil.ExecError, e:
+ raise self.Error("can't get public key for %s: " % self.path + str(e))
+
+ @property
+ def fingerprint(self):
+ return hashlib.sha1(self.public).hexdigest()
+
+class TempPrivateKey(TempFile, PrivateKey):
+ def __init__(self):
+ TempFile.__init__(self, prefix='key_')
+ os.remove(self.path)
+
+ executil.getoutput("ssh-keygen -N '' -f %s" % self.path)
+ os.remove(self.path + ".pub")
+
+ PrivateKey.__init__(self, self.path)
+
class SSH:
class Error(Exception):
pass
@@ -95,31 +130,26 @@ def command(self, command, pty=False):
callback=self.callback,
pty=pty)
- def copy_id(self, key_path):
- if not key_path.endswith(".pub"):
- key_path += ".pub"
+ def copy_id(self, key):
+ if not isinstance(key, PrivateKey):
+ key = PrivateKey(key)
command = 'mkdir -p $HOME/.ssh; cat >> $HOME/.ssh/authorized_keys'
command = self.command(command)
- command.tochild.write(file(key_path).read())
+ command.tochild.write("%s %s\n" % (key.public, key.fingerprint))
command.tochild.close()
try:
command.close()
except command.Error, e:
raise self.Error("can't add id to authorized keys: " + str(e))
- def remove_id(self, key_path):
- if not key_path.endswith(".pub"):
- key_path += ".pub"
-
- vals = file(key_path).read().split()
- if not vals[0].startswith('ssh'):
- raise self.Error("invalid public key in " + key_path)
- id = vals[-1]
+ def remove_id(self, key):
+ if not isinstance(key, PrivateKey):
+ key = PrivateKey(key)
- command = 'sed -i "/%s/d" $HOME/.ssh/authorized_keys' % id
+ command = 'sed -i "/%s/d" $HOME/.ssh/authorized_keys' % key.fingerprint
command = self.command(command)
try:
View
37 cloudtask/task.py
@@ -24,6 +24,7 @@
--force Don't ask for confirmation
+ --ssh-identity= SSH identity keyfile to use (defaults to ~/.ssh/identity)
--hub-apikey= Hub API KEY (required if launching workers)
--snapshot-id= Launch instance from a snapshot ID
@@ -87,27 +88,7 @@
from reporter import Reporter
from watchdog import Watchdog
-from temp import TempFile
-import executil
-import uuid
-
-class TempSSHKey(TempFile):
- def __init__(self):
- TempFile.__init__(self, prefix='key_')
- os.remove(self.path)
-
- self.uuid = uuid.uuid4()
- executil.getoutput("ssh-keygen -N '' -f %s -C %s" % (self.path, self.uuid))
-
- @property
- def public(self):
- return self.path + ".pub"
-
- def __del__(self):
- if os.getpid() == self.pid:
- os.remove(self.public)
-
- TempFile.__del__(self)
+import ssh
class Task:
@@ -164,6 +145,7 @@ def filter(job):
answer = raw_input("Is this really what you want? [yes/no] ")
if answer:
break
+
sys.stdin = orig_stdin
if answer.lower() != "yes":
@@ -261,6 +243,12 @@ def main(cls):
taskconf.overlay = abspath(val)
+ elif opt == '--ssh-identity':
+ if not isfile(val):
+ error("ssh-identity '%s' not a file" % val)
+
+ taskconf.ssh_identity = abspath(val)
+
elif opt == '--split':
taskconf.split = int(val)
if taskconf.split < 1:
@@ -376,6 +364,11 @@ def main(cls):
@classmethod
def work(cls, jobs, session, taskconf):
+ if taskconf.ssh_identity:
+ sshkey = ssh.PrivateKey(taskconf.ssh_identity)
+ else:
+ sshkey = ssh.TempPrivateKey()
+
def status(msg):
timestamp = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
session.logs.manager.write("%s :: session %d %s\n" % (timestamp, session.id, msg))
@@ -403,8 +396,6 @@ def terminate(sig, f):
executor = None
work_started = time.time()
- sshkey = TempSSHKey()
-
try:
executor = CloudExecutor(session.logs, taskconf, sshkey)
for job in jobs:
View
4 cloudtask/taskconf.py
@@ -32,6 +32,8 @@ class TaskConf:
ami_id = None
snapshot_id = None
+ ssh_identity = None
+
ec2_region = 'us-east-1'
ec2_size = 'm1.small'
ec2_type = 's3'
@@ -80,7 +82,7 @@ def fmt(self):
sio = StringIO()
table = []
- for attr in ('split', 'command', 'hub-apikey',
+ for attr in ('split', 'command', 'ssh-identity', 'hub-apikey',
'ec2-region', 'ec2-size', 'ec2-type',
'user', 'backup-id', 'ami-id', 'snapshot-id', 'workers',
'overlay', 'post', 'pre', 'timeout', 'report'):

0 comments on commit 971b282

Please sign in to comment.
Something went wrong with that request. Please try again.