# $Id: sftp.py,v 1.1 2010-05-28 18:42:29 wirawan Exp $
#
# wpylib.net.sftp_robot module
# Created: 20100226
# Wirawan Purwanto
#
"""
SFTP-related module.

A better implementation is based on paramiko secure shell library
(http://www.lag.net/paramiko/).
A clunky implementation is based on `sftp' program. But the status inquiry
only works for Linux OS since it uses /proc/$PID approach.
"""

import getpass
import os
import stat
import sys
import time

try:
  import paramiko
  has_paramiko = True
except:
  has_paramiko = False

if has_paramiko:
  class sftp_client(object):
    """Python-controllable SFTP client.
    It's better to have no password (i.e. using private/public key authentication)
    to avoid additional security issues with python, password being passed around
    in clear text."""
    def __init__(self, hostname, username=None, port=22, password=None):
      conn = paramiko.SSHClient()
      if username == None:
        username = getpass.getuser()
  
      self.conn = conn
      self.username = username
      self.hostname = hostname
      self.port = hostname
  
      # Default init stage:
      conn.load_system_host_keys()
      #print "here1"
      conn.load_host_keys(os.path.expanduser("~/.ssh/known_hosts"))
      #print "here1"
      conn.set_missing_host_key_policy(paramiko.RejectPolicy()) # default=paranoid
      #print "here1"
      conn.connect(hostname=hostname, username=username, password=password,
                   port=port,
                  )
      sftp = conn.open_sftp()
      self.sftp = sftp
  
    def cd(self, path):
      self.sftp.chdir(path)
  
    def getcwd(self):
      self.sftp.getcwd()
  
    def get(self, remotepath, localpath, callback=None):
      # TODO: progress bar
      self.sftp.get(remotepath, localpath, callback)
      # Preserve the attributes as much as I can do it
      stats = self.sftp.stat(remotepath)
      os.utime(localpath, (stats.st_atime, stats.st_mtime))
      os.chmod(localpath, stats.st_mode)

    def put(self, localpath, remotepath, callback=None):
      # TODO: progress bar
      self.sftp.put(localpath, remotepath, callback)
      # Preserve the attributes as much as I can do it
      stats = os.stat(localpath)
      self.sftp.utime(remotepath, (stats.st_atime, stats.st_mtime))
      self.sftp.chmod(remotepath, stats.st_mode)

    # os.path-like functions

    def _is_test(self, path, test_proc):
      # Credit:
      # http://stackoverflow.com/questions/6674862/recursive-directory-download-with-paramiko
      #from stat import S_ISDIR
      try:
        return test_proc(self.sftp.stat(path).st_mode)
      except IOError:
        return False

    def isdir(self, path):
      return self._is_test(path, stat.S_ISDIR)

    def islink(self, path):
      """This function is not reliable, I found.
      In my test cases, softlinks were not be detected correctly.
      This might have to do with softlink support in the SFTP server,
      but most likely it was the limitation of paramiko itself.
      CAVEAT EMPTOR.
      """
      return self._is_test(path, stat.S_ISLNK)

    def isfile(self, path):
      return self._is_test(path, stat.S_ISREG)

    def close(self):
      self.sftp.close()


# TODO: recode the sftp driver from

class sftp_driver(object):
  """Python controllable SFTP client built upon OpenSSH's sftp program."""
  def __init__(self, hostname, username=None, port=22, password=None):
    raise NotImplementedError
  # TODO


class sftp_status_callback(object):
  """SFTP progress bar for sftp_client object above.

  Example usage:

      statproc = sftp_status_callback(fname=remote_root + "/myfile.txt")
      sftp.get(remote_root + "/myfile.txt", "mylocalfile.txt", statproc)
      statproc.done()

  """
  def __init__(self, fname, out=sys.stderr):
    self.first = True
    self.fname = fname
    self.out = out
    self.marks = [ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90 ]
    self.nbars = 25
    self.lastbar = 0
    self.begin_time = time.time()
  def printbar(self, bar):
    if bar > self.lastbar:
      dbar = bar - self.lastbar
      self.out.write("*" * dbar)
      self.out.flush()
      self.lastbar = bar
  def __call__(self, nbytes, filesize):
    if self.first:
      self.first = False
      self.filesize = filesize
      self.last = 0
      self.out.write("%-40s : %12d " % (self.fname[-40:], self.filesize))
      self.out.flush()
    self.cursize = nbytes
    curbar = int(float(nbytes) / self.filesize * self.nbars)
    self.printbar(curbar)
  def done(self):
    self.end_time = time.time()
    dtime = self.end_time - self.begin_time
    self.out.write("  %.2f kibps\n" % (self.filesize / 1024.0 / dtime))
    self.out.flush()