daily pastebin goal
66%
SHARE
TWEET

Untitled

a guest Sep 23rd, 2016 59 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. from contextlib import contextmanager
  2. from paramiko import AuthenticationException
  3. import pysftp
  4.  
  5. @contextmanager
  6. def open_sftp(link, login, mode='r'):
  7.     """Context manager to read/write a file via SFTP
  8.  
  9.     :param link: Path to file
  10.     :param login: SFTP username/password info in the format
  11.  
  12.                     {
  13.                         'username': "Username",
  14.                         'password': "p@ssw0rd"
  15.                     }
  16.  
  17.     :param mode: The mode to open the file with. Either 'r' or 'w'
  18.     :returns: File object for the specified file
  19.     :raises:
  20.         :ValueError: If there is an error with the input
  21.         :AuthenticationException: If could not connect
  22.         :IOError: If file cannot be found/opened
  23.     """
  24.     username = login.get('username')
  25.     password = login.get('password')
  26.  
  27.     # Parse link into host, directory, and file
  28.     # Not much error checking, depends on link matching:
  29.     #   [protocol]://[host]/[path (optional)]/[file]
  30.  
  31.     # Remove the sftp protocol from the link string
  32.     link = link.replace(u"sftp://", u"")
  33.  
  34.     link_split = link.split('/')
  35.     if len(link_split) < 2:
  36.         raise ValueError("Error parsing link")
  37.     link_file = link_split.pop()
  38.     link_path = link_split
  39.     if not link_file:
  40.         raise ValueError("Error parsing link: No file path")
  41.  
  42.     cinfo = {
  43.         'host': link_split.pop(0),
  44.         'username': username,
  45.         'password': password
  46.     }
  47.     if not cinfo.get('host'):
  48.         raise ValueError("Error parsing link: No host")
  49.     if u":" in cinfo['host']:
  50.         cinfo['host'], cinfo['port'] = cinfo['host'].split(u':')
  51.         cinfo['port'] = int(cinfo['port'])
  52.  
  53.     # Connect to host. Sometimes it takes a couple tries... =(
  54.     sftp = None
  55.     for i in xrange(5):
  56.         try:
  57.             sftp = pysftp.Connection(**cinfo)
  58.             break
  59.         except AuthenticationException:
  60.             warnings.warn("Could not authenticate... Will re-attempt")
  61.             import time as t
  62.             t.sleep(5)
  63.  
  64.     # Try one more time... if this fails, raise the AuthenticationError
  65.     if not sftp:
  66.         sftp = pysftp.Connection(**cinfo)
  67.  
  68.     try:
  69.         for directory in link_path:
  70.             try:
  71.                 sftp.chdir(directory)
  72.             except IOError as exc:
  73.                 raise IOError("Could not navigate to file: {}".format(exc))
  74.  
  75.         try:
  76.             file_obj = sftp.open(link_file, mode=mode)
  77.         except IOError as exc:
  78.             raise IOError("Could not open file: {}".format(exc))
  79.  
  80.         yield file_obj
  81.     finally:
  82.         sftp.close()
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top