# Implementation of a subset of 'ssh' based on fabric and paramiko
# This is used for remote development when the user selects "builtin"
# SSH Implementation.  It is done this way (in a separate process
# and not integrated into Wing's process space) because fabric and
# paramiko spawn threads that conflict with Wing's totally async
# design and also some of Qt's expectations, leading to threading
# deadlocks.

import os
import sys
import getpass
import traceback

# Need to add the fabric runtime site-packages directory to the Python path
# since this isn't the same as the main site-packages directory
pyver = '{}.{}'.format(sys.version_info.major, sys.version_info.minor)
pyruntime_dirname = 'runtime-python{}'.format(pyver)
runtimedir = sys.executable
runtimedir = runtimedir[:runtimedir.find('runtime-python{}'.format(pyver))]
if sys.platform != 'win32':
  sitedir = 'runtime-pip-fabric-py{}/lib/python{}/site-packages'.format(pyver, pyver)
  sitedir = os.path.join(runtimedir, sitedir)
else:
  sitedir = r'runtime-pip-fabric-py{}\lib\site-packages'.format(pyver)
sitedir = os.path.join(runtimedir, sitedir)
sys.path.append(sitedir)

from fabric import Connection
import paramiko
import invoke

kSuccess = 0
kAuthenticationFailed = 2
kCopyFailed = 3
kBadHostKey = 4

def pw_prompt(prompt):
  
  if not prompt.endswith(' '):
    prompt += ' '
  
  # Can't use getpass on Windows when this is called from Wing b/c there is no
  # real pseudo TTY implementation on Wing and the output from getpass is
  # screwy and messes up interaction.  Doesn't matter b/c the password is not 
  # echoed in that case anyway and this tool is not really designed to be used 
  # from the command line.  If it were then should add code here to detect when 
  # in a Windows terminal vs. not.
  if sys.platform != 'win32':
    return getpass.getpass(prompt)
  else:
    print("WARNING: Password will be echoed if running in a terminal")
    return input(prompt)
  
def find_key_files():
    
  found_files = []
  
  if sys.platform == 'win32':
    dirs = ('.ssh', 'ssh')
  else:
    dirs = ('.ssh', )

  # Matches what Paramiko and (I think) OpenSSH does so user is not surprized
  # by which keys they are being prompted for
  keytypes = ('rsa', 'dsa', 'ecdsa', 'ed25519')
  
  for directory in dirs:
    for keytype in keytypes:
      full_path = "~/{}/id_{}".format(directory, keytype)
      if os.path.isfile(os.path.expanduser(full_path)):
        found_files.append(full_path)

  return found_files

class CConfig:
  
  def __init__(self):
    
    self.port = 22
    self.keyfile = None
    self.connect_timeout = None
    self.use_agent = True
    self.strict_host_key = True
    self.verbose = False
    self.forward_agent = False
    self.forward_x11 = False
    self.reverse_tunnel = None
    self.port_forwards = []
    self.hostname = None
    self.copy_operation = False
    self.clear_all_forwardings = False
    
  def ParseArgv(self, args):
    
    i = 0
    while i < len(args):
      
      arg = args[i]
      
      if arg.lower() == '-p':
        self.port = int(args[i+1])
        i += 1
        
      elif arg == '-i':
        self.keyfile = args[i+1]
        i += 1
        
      elif arg == '-o':
        
        if args[i+1].lower().startswith('stricthostkeychecking='):
          parts = args[i+1].split('=')
          if parts[1].lower() in ('true', 'yes', '1'):
            self.strict_host_key = True
          else:
            self.strict_host_key = False
            
        elif args[i+1].lower().startswith('connecttimeout='):
          parts = args[i+1].split('=')
          self.connect_timeout = float(parts[1])
        
        elif args[i+1].lower() == 'identityagent=none':
          self.use_agent = False
        
        elif args[i+1].lower().startswith('clearallforwardings='):
          parts = args[i+1].split('=')
          if parts[1].lower() in ('true', 'yes', '1'):
            self.clear_all_forwardings = True
          else:
            self.clear_all_forwardings = False
        
        elif args[i+1].lower().startswith('forwardx11='):
          parts = args[i+1].split('=')
          if parts[1].lower() in ('true', 'yes', '1'):
            self.forward_x11 = True
          else:
            self.forward_x11 = False
        
        elif args[i+1].lower().startswith('forwardagent='):
          parts = args[i+1].split('=')
          if parts[1].lower() in ('true', 'yes', '1'):
            self.forward_agent = True
          else:
            self.forward_agent = False
        
        else:
          raise RuntimeError("Unknown -o option: {}".format(args[i+1]))

        i += 1
        
      elif arg == '-v':
        self.verbose = True
        
      # XXX X11 forwarding w/ fabric is actually not implemented and disabled in the 
      # XXX UI when fabric is being used as the SSH implementation
      elif arg in ('-X', '-Y'):
        self.forward_x11 = True
      elif arg == '-x':
        self.forward_x11 = False
        
      elif arg == '-R':
        remote_port, hostname, local_port = args[i+1].split(':')
        remote_port = int(remote_port)
        local_port = int(local_port)
        self.reverse_tunnel = (remote_port, local_port)
        i += 1
        
      elif arg == '-L':
        local_host, local_port, remote_host, remote_port = args[i+1].split(':')
        local_port = int(local_port)
        remote_port = int(remote_port)
        self.port_forwards.append((local_port, remote_port))
        i += 1
        
      elif arg == '--copy':
        source = args[i+1]
        dest = args[i+2]
        self.copy_operation = [source, dest]
        i += 2
      
      elif self.hostname is None:
        self.hostname = args[i]
        if '@' not in self.hostname:
          self.hostname = getpass.getuser() + '@' + self.hostname
          
      else:
        self.cmd = args[i:]
        break
        
      i += 1    
      
    return config
  
class CMissingHostKeyMgr(paramiko.client.MissingHostKeyPolicy):
  
  def __init__(self, owner):
    self.owner = owner
    
  def missing_host_key(self, client, hostname, key):
    print("Found unknown host key for {}: {}".format(hostname, key.get_base64()))
    accept = input("Are you sure you want to continue (y/n)? ")
    if accept.strip().lower() in ('y', 'yes'):
      if isinstance(key, paramiko.dsskey.DSSKey):
        keytype = 'ssh-dsa'
      elif isinstance(key, paramiko.rsakey.RSAKey):
        keytype = 'ssh-rsa'
      elif isinstance(key, paramiko.ecdsakey.ECDSAKey):
        keytype = 'ecdsa'
      elif isinstance(key, paramiko.ed25519key.Ed25519Key):
        keytype = 'ed25519'
      else:
        print("Unsupported key type: {} - cannot store key".format(type(key)))
      self.owner.conn.client.get_host_keys().add(hostname, keytype, key)
      self.owner.conn.client.save_host_keys(self.owner.hostkeys)
      return
    else:
      raise RuntimeError("Rejected host key")
  
class CFabricConnection:
  
  def __init__(self, config:CConfig):
    
    self.conn = None
    self.result = None
    self.sftp = None
    self.output_pos = 0
    self.output_trunc_pos = 0
    self.all_output = ''

    self.entered_passphrases = {}
    self.entered_password = False
    
    self.hostkeys = os.path.expanduser('~/.ssh/known_hosts')
    self.c = config
    
    self.connected = False

  def Connect(self):
    """Try to connect the remote host using authentication in this order
    (which matches what OpenSSH seems to do; the first two steps are
    reversed in Paramiko):
    
       - Any key found in an SSH agent, if not disabled by user preference
       - User-specified key file, if any.
       - If no user-specified key was given, any discoverable key in ~/.ssh
         (and ~/ssh on Windows)
       - Login password authentication
    
    This replicates what is done in Paramiko's client.SSHClient._auth and
    is done here so we know which private key we are attempting to access
    when prompting for passphrases.
    
    """
    
    # Tries each possibility.  When TryToConnect succeeds, it does not return
    # and eventually calls sys.exit() when the connection is dropped.
    
    print("Connecting to {}".format(self.c.hostname))

    # Attempt SSH agent unless disabled
    if self.c.use_agent:
      print("Trying SSH user agent")
      self.TryToConnect(None, allow_agent=True)
      
    # Attempt user-specified key file, if any
    if self.c.keyfile is not None:
      if os.path.exists(os.path.expanduser(self.c.keyfile)):
        print("Trying user-specified key", self.c.keyfile)
        self.TryToConnect(self.c.keyfile)
      else:
        print("Warning: Key file {} does not exist".format(self.c.keyfile))
      
    # Attempt findable keys if user didn't provide one
    if self.c.keyfile is None:
      for keyfile in find_key_files():
        print("Trying found key", keyfile)
        self.TryToConnect(keyfile)
      
    # Use login password authentication
    print("Falling back on login password")
    self._CollectPassword()
      
  def TryToConnect(self, key_filename, allow_agent=False, passphrase=None, password=None):    

    connect_kwargs = {
      'look_for_keys': False,
      'allow_agent': allow_agent,
    }    
    if key_filename:
      connect_kwargs['key_filename'] = os.path.expanduser(key_filename)
    if passphrase:
      connect_kwargs['passphrase'] = passphrase
    if password:
      connect_kwargs['password'] = password
    
    tunnels = []
    
    try:
      self._CreateSshDir()
      
      self.conn = Connection(self.c.hostname, port=self.c.port, connect_kwargs=connect_kwargs,
                             forward_agent=self.c.forward_agent)
      
      # Set up missing host key handling
      try:
        self.conn.client.load_host_keys(self.hostkeys)
      except IOError:
        self.Log("Could not read host keys from {}".format(self.hostkeys))
      self.missing_host_key_mgr = CMissingHostKeyMgr(self)
      self.conn.client.set_missing_host_key_policy(self.missing_host_key_mgr)
      
      # Open connect to remote host explicitly (Fabric does this implicitly 
      # as needed but as Tim says "explicit is better than implicit")
      self.conn.open()
      
      # Start up reverse tunnel, if any.  Weird implementation is due to the API
      # assuming you're going to be calling things in a 'with' -- basically assuming
      # this is used for short-lived calls that aren't really async (the async support
      # in Fabric is basically not real)
      if self.c.reverse_tunnel:
        rgen = self.conn.forward_remote(*self.c.reverse_tunnel)
        tunnels.append(rgen)
        next(rgen.gen)
        self.Log("Port {} forwarded to local address 127.0.0.1 port {}".format(*self.c.reverse_tunnel))

      # Start up any user-defined forward tunnels; same issues as above
      for local, remote in self.c.port_forwards:
        fgen = self.conn.forward_local(local, remote)
        tunnels.append(fgen)
        next(fgen.gen)
        self.Log("Port forward success local {} to remote {}".format(local, remote))
        
      if self.c.copy_operation:
        self.sftp = self.conn.sftp()
        self.RecursiveCopy(*self.c.copy_operation)
      else:
        # The logic for 'warn' appears to be backward in the API - it warns about
        # there being no stderr for a pty (by raising a fatal exception!) unless
        # we send in warn=True
        cmdline = self.ArgsToCmdLine(self.c.cmd)
        print("Running: {}".format(cmdline))
        self.result = self.conn.run(cmdline, pty=True, warn=True)
      
    # Bad host key (mismatch)
    except paramiko.ssh_exception.BadHostKeyException as e:
      hostname = self.c.hostname.split('@')[-1]
      print("WARNING: Bad host key received for host {} -- does not match the host key".format(hostname))
      print('WARNING: stored in {}.  This may indicate a "man in the middle" attack!'.format(self.hostkeys))
      print("Details:", str(e))
      print("Connection canceled")
      sys.exit(kBadHostKey)
      
    # Need a passphrase to unlock private key
    except paramiko.PasswordRequiredException:
      if key_filename is not None:
        self._CollectPassphrase(key_filename)
      return
  
    # Rejected key or incorrectly entered password
    except paramiko.AuthenticationException as e:
      if key_filename is not None:
        if self.c.verbose:
          print("Host {} rejected key {} -- {}".format(self.c.hostname, key_filename, e.args[0]))
      return 
    
    # Connection failed for some other reason
    except paramiko.SSHException as e:
      if key_filename is not None:
        self._CollectPassphrase(key_filename)
        return
      elif 'no authentication' in e.args[0].lower():
        if not allow_agent:
          print("Connection failed: {}".format(e.args[0]))
        return
      else:
        print("Unexpected connection failure: {}".format(e.args[0]))
        traceback.print_exc()
        return

    # Close down any tunnels
    for gen in tunnels:
      try:
        next(gen.gen)
      except StopIteration:
        pass
      
    # Close the connection
    self.conn.close()
    
    # If this was a copy operation, then it's done if reach here
    if self.sftp is not None:
      sys.exit(kSuccess)
      return
    
    if self.result is None:
      sys.exit(kAuthenticationFailed)
    else:
      sys.exit(self.result.exited)
      
  def _CollectPassphrase(self, key_filename):
    
    if self.entered_passphrases.get(key_filename):
      print("Invalid passphrase for key {}".format(key_filename))
      return
      
    self.entered_passphrases[key_filename] = True
    
    expanded_keyfile = os.path.expanduser(key_filename)

    display_keyfile = (key_filename if sys.platform != 'win32' 
                       else expanded_keyfile.replace('/', '\\'))
    pp = pw_prompt("Enter passphrase for keyfile {}:".format(display_keyfile))
    
    self.TryToConnect(key_filename, passphrase=pp)
    
  def _CollectPassword(self):

    if self.entered_password:
      print("Invalid password - connection failed")
      sys.exit(2)
      return
      
    self.entered_password = True
      
    print("Enter login password for {}".format(self.c.hostname))
    pp = pw_prompt("Password:")
    
    self.TryToConnect(None, password=pp)
    
  def _CreateSshDir(self):
    """ Create ~/.ssh if needed """
      
    # Make sure an ssh dir exists.  We use ~/ssh on Windows if it exists
    # but ~/.ssh does not.  If neither exist, we create ~/.ssh
    ssh_dir = os.path.expanduser('~/.ssh')
    if not os.path.exists(ssh_dir):
      if sys.platform == 'win32' and os.path.exists(os.path.expanduser('~/ssh')):
        ssh_dir = os.path.expanduser('~/ssh')
      else:
        os.makedirs(ssh_dir, mode=0o700)
    
    # Paramiko assumes ~/.ssh/known_hosts exists
    known_hosts_fn = os.path.join(ssh_dir, 'known_hosts')
    if not os.path.exists(known_hosts_fn):
      f = open(known_hosts_fn, 'wb')
      f.close()
      
  def Log(self, *args):
    
    if self.c.verbose:
      print(*args)
      
  def RecursiveCopy(self, local_dir, remote_dir):
    """Copy the given local directory recursively to the given remote directory,
    creating any directories as needed"""
    
    local_dir = os.path.expanduser(local_dir)
    
    remote_dir = self.sftp.normalize(remote_dir)
    remote_dir = os.path.join(remote_dir, os.path.basename(local_dir))
    
    print("Copy {} -> {}".format(local_dir, remote_dir))
    
    for lparent, ldirs, lfiles in os.walk(local_dir):
      relparent = lparent[len(local_dir)+1:]
      if relparent:
        rdir = os.path.join(remote_dir, relparent)
      else:
        rdir = remote_dir
      if sys.platform == 'win32':
        rdir = rdir.replace('\\', '/')
        
      try:
        self.Log("mkdir", rdir)
        self.sftp.mkdir(rdir)
      except IOError as e:
        # errno is None if it failed because the directory already exists
        if e.errno:
          self.Log("Failed to create remote directory (err={}): {}".format(e.errno, rdir))
          traceback.print_exc()
          sys.exit(kCopyFailed)
  
      for file in lfiles:
        local_fp = os.path.join(lparent, file)
        remote_fp = os.path.join(remote_dir, relparent, file)
        if sys.platform == 'win32':
          remote_fp = remote_fp.replace('\\', '/')
        self.Log("File {} -> {}".format(local_fp, remote_fp))
        self.sftp.put(local_fp, remote_fp)
    
  def ArgsToCmdLine(self, argv):
    
    args2=[]
    for arg in argv:
      if arg == '':
        arg2 = '""'
      elif ' ' not in arg and '\t' not in arg:
        arg2 = arg
      elif arg[0] == arg[-1] and arg[0] == '"':
        arg2 = arg
      else:
        arg2 = ('"' + arg.replace('"', '\\"') + '"')
      args2.append(arg2)
        
    # Special case single arg with spaces to handle how the CReverseSSHTunnel works
    # This is needed b/c of how 'invoke' works (or doesn't, due to using command
    # line and not list for args -- same design flaw we had in Wing for a long time)
    if len(args2) == 1 and args2[0][0] == args2[0][-1] and args2[0][0] == '"':
      args2 = [args2[0][1:-1]]
      
    return ' '.join(args2)

if __name__ == '__main__':
  
  config = CConfig()
  config.ParseArgv(sys.argv[1:])
  
  proc = CFabricConnection(config)
  proc.Connect()
  
  
  