python/cumulus/store/sftp.py

   1 # vim: ai ts=4 sts=4 et sw=4
   2 from paramiko import Transport, SFTPClient, RSAKey, DSSKey
   3 from paramiko.config import SSHConfig
   4 import paramiko.util
   5 from cumulus.store import Store, type_patterns, NotFoundError
   6 import os, os.path
   7 import getpass
   8 import re
   9 import sys
  10 #needed for python 2.5
  11 try:
  12     from __future__ import with_statement
  13 except ImportError, e:
  14     raise ImportError('We need the "with" statement, avalible with python >=2.5')
  15
  16
  17 class SSHHostConfig(dict):
  18     def __init__(self, hostname, user = None, filename = None):
  19         dict.__init__()
  20         #set defaults
  21         if filename == None:
  22             filename = os.path.expanduser('~/.ssh/config')
  23         self['port'] = 22
  24         self['user'] = getpass.getuser()
  25         self['hostname'] = hostname
  26         self['hostkeyalias'] = hostname
  27
  28         #read config file
  29         ssh_config = SSHConfig()
  30         with open(filename) as config_file:
  31             ssh_config.parse(config_file)
  32         self.update(ssh_config.lookup(hostname))
  33
  34         if user != None:
  35             self['user'] = user
  36
  37 class SFTPStore(Store):
  38     """implements the sftp:// storage backend
  39
  40         configuration via openssh/sftp style urls and
  41         .ssh/config files
  42
  43         does not support password authentication or password
  44         protected authentication keys"""
  45     def __init__(self, url, **kw):
  46         if self.path.find('@') != -1:
  47             user, self.path = self.path.split('@')
  48         else:
  49             user = None
  50
  51         if self.path.find(':') != -1:
  52             host, self.path = selfpath.split(':')
  53         else:
  54             host, self.path = self.path.split('/', 1)
  55
  56         self.config = SSHHostConfig(host, user)
  57
  58
  59         host_keys = paramiko.util.load_host_keys(os.path.expanduser('~/.ssh/known_hosts'))
  60         self.hostkey = host_keys[config['hostkeyalias']].values()[0]
  61
  62         if(config.has_key('identityfile')):
  63             key_file = os.path.expanduser(host_config['identityfile'])
  64             #not really nice but i don't see a cleaner way atm...
  65             try:
  66                 self.auth_key = RSAKey (filename = key_file)
  67             except SSHException, e:
  68                 if e.message == 'Unable to parse file':
  69                     self.auth_key = DSAKey (filename = key_file)
  70                 else:
  71                     raise
  72         else:
  73             filename = os.path.expanduser('~/.ssh/id_rsa')
  74             if os.path.exists(filename):
  75                 self.auth_key = RSAKey(filename = filename)
  76             else:
  77                 filename = os.path.expanduser('~/.ssh/id_dsa')
  78                 if (os.path.exists(filename)):
  79                     self.auth_key = DSSKey (filename = filename)
  80
  81         self.__connect()
  82
  83     def __connect(self):
  84         self.t = Transport((self.config['hostname'], self.config['port']))
  85         self.t.connect(username = self.config['user'], pkey = self.auth_key)
  86         self.client = SFTPClient.from_transport(self.t)
  87         self.client.chdir(self.path)
  88
  89     def __build_fn(self, name):
  90         return "%s/%s" % (self.path,  name)
  91
  92     def list(self, type):
  93         return filter(type_patterns[type].match, self.client.listdir(self.path))
  94
  95     def get(self, type, name):
  96         return self.client.open(filename = self.__build_fn(name), mode = 'rb')
  97
  98     def put(self, type, name, fp):
  99         remote_file = self.client.open(filename = self.__build_fn(name), mode = 'wb')
 100         buf = fp.read(4096)
 101         while (len(buf) > 0):
 102             remote_file.write(buf)
 103             buf = fp.read(4096)
 104         remote_file.close()
 105
 106     def delete(self, type, name):
 107         self.client.remove(filename = self.__build_fn(name))
 108
 109     def stat(self, type, name):
 110         stat = self.client.stat(filename = self.__build_fn(name))
 111         return {'size': stat.st_size}
 112
 113     def close(self):
 114         """connection has to be explicitly closed, otherwise
 115             it will hold the process running idefinitly"""
 116         self.client.close()
 117         self.t.close()
 118
 119 Store = SFTPStore