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