/*
 * This file is part of the KFTPGrabber project
 *
 * Copyright (C) 2003-2006 by the KFTPGrabber developers
 * Copyright (C) 2003-2006 Jernej Kos <kostko@jweb-network.net>
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation; either version 2 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * is provided AS IS, WITHOUT ANY WARRANTY; without even the implied
 * warranty of MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, and
 * NON-INFRINGEMENT.  See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Steet, Fifth Floor, Boston,
 * MA 02110-1301, USA.
 *
 * In addition, as a special exception, the copyright holders give
 * permission to link the code of portions of this program with the
 * OpenSSL library under certain conditions as described in each
 * individual source file, and distribute linked combinations
 * including the two.
 *
 * You must obey the GNU General Public License in all respects
 * for all of the code used other than OpenSSL.  If you modify
 * file(s) with this exception, you may extend this exception to your
 * version of the file(s), but you are not obligated to do so.  If you
 * do not wish to do so, delete this exception statement from your
 * version.  If you delete this exception statement from all source
 * files in the program, then also delete it here.
 */

#include "sftpsocket.h"
#include "cache.h"
#include "misc/kftpconfig.h"

#include <ntqdir.h>

#include <tdelocale.h>
#include <kstandarddirs.h>
#include <tdeio/job.h>
#include <tdeio/renamedlg.h>

#include <sys/stat.h>
#include <fcntl.h>

namespace KFTPEngine {

SftpSocket::SftpSocket(Thread *thread)
  : Socket(thread, "sftp"),
    m_login(false)
{
}

SftpSocket::~SftpSocket()
{
}

int addPermInt(int &x, int n, int add)
{
  if (x >= n) {
    x -= n;
    return add;
  } else {
    return 0;
  }
}

int SftpSocket::intToPosix(int permissions)
{
  int posix = 0;
  TQString str = TQString::number(permissions);
  
  int user = str.mid(0, 1).toInt();
  int group = str.mid(1, 1).toInt();
  int other = str.mid(2, 1).toInt();
  
  posix |= addPermInt(user, 4, S_IRUSR);
  posix |= addPermInt(user, 2, S_IWUSR);
  posix |= addPermInt(user, 1, S_IXUSR);
  
  posix |= addPermInt(group, 4, S_IRGRP);
  posix |= addPermInt(group, 2, S_IWGRP);
  posix |= addPermInt(group, 1, S_IXGRP);
  
  posix |= addPermInt(other, 4, S_IROTH);
  posix |= addPermInt(other, 2, S_IWOTH);
  posix |= addPermInt(other, 1, S_IXOTH);
  
  return posix;
}


// *******************************************************************************************
// ***************************************** CONNECT *****************************************
// *******************************************************************************************

class SftpCommandConnect : public Commands::Base {
public:
    enum State {
      None,
      ConnectComplete,
      LoginComplete
    };
    
    ENGINE_STANDARD_COMMAND_CONSTRUCTOR(SftpCommandConnect, SftpSocket, CmdConnect)
    
    void process()
    {
      KURL url = socket()->getCurrentUrl();
      
      switch (currentState) {
        case None: {
          // Set connection info
          SSH_OPTIONS *sshOptions = options_new();
          options_set_username(sshOptions, (char*) url.user().ascii());
          options_set_host(sshOptions, url.host().ascii());
          options_set_port(sshOptions, url.port());
          options_set_timeout(sshOptions, 10, 0);
          
          socket()->m_sftpSession = 0;
          socket()->m_sshSession = ssh_connect(sshOptions);
          
          if (!socket()->sshSession()) {
            socket()->emitEvent(Event::EventMessage, i18n("Unable to establish SSH connection (%1)").arg(ssh_get_error(0)));
            socket()->emitError(ConnectFailed);
            return;
          }
          
          socket()->emitEvent(Event::EventState, i18n("Logging in..."));
          socket()->emitEvent(Event::EventMessage, i18n("Connected with server, attempting to login..."));
          
          currentState = ConnectComplete;
        }
        case ConnectComplete: {
          SSH_SESSION *sshSession = socket()->sshSession();
          TQString password;
          
          // Check if a public key password was supplied using the wakeup event
          if (isWakeup()) {
            PubkeyWakeupEvent *event = static_cast<PubkeyWakeupEvent*>(m_wakeupEvent);
            password = event->password;
          }
          
          // Try the public key auth with the set password (if any)
          int pkey_ret = ssh_userauth_autopubkey(sshSession, (char*) password.ascii());
          if (pkey_ret == -666) {
            // Make a password request
            socket()->emitEvent(Event::EventPubkeyPassword);
            return;
          } else if (pkey_ret != SSH_AUTH_SUCCESS) {
            // First let's try the keyboard-interactive authentification
            if (keyboardInteractiveLogin() != SSH_AUTH_SUCCESS) {
              // If this fails, let's try the password authentification
              if (ssh_userauth_password(sshSession, NULL, (char*) url.pass().ascii()) != SSH_AUTH_SUCCESS) {
                socket()->emitEvent(Event::EventMessage, i18n("Login has failed."));
                socket()->emitError(LoginFailed);
                
                socket()->protoAbort();
                return;
              }
            } else {
              socket()->emitEvent(Event::EventMessage, i18n("Keyboard-interactive authentication succeeded."));
            }
          } else {
            socket()->emitEvent(Event::EventMessage, i18n("Public key authentication succeeded."));
          }
          
          currentState = LoginComplete;
        }
        case LoginComplete: {
          socket()->m_sftpSession = sftp_new(socket()->sshSession());
          
          if (!socket()->sftpSession()) {
            socket()->emitEvent(Event::EventMessage, i18n("Unable to initialize SFTP channel."));
            socket()->emitError(LoginFailed);
            
            socket()->protoAbort();
            return;
          }
          
          if (sftp_init(socket()->sftpSession())) {
            socket()->emitEvent(Event::EventMessage, i18n("Unable to initialize SFTP."));
            socket()->emitError(LoginFailed);
            
            socket()->protoAbort();
            return;
          }
          
          // Get the current directory
          char *cwd = sftp_canonicalize_path(socket()->sftpSession(), "./");
          socket()->setDefaultDirectory(socket()->remoteEncoding()->decode(cwd));
          socket()->setCurrentDirectory(socket()->remoteEncoding()->decode(cwd));
          delete cwd;
          
          socket()->emitEvent(Event::EventMessage, i18n("Connected."));
          socket()->emitEvent(Event::EventConnect);
          socket()->m_login = true;
          
          socket()->resetCommandClass();
          break;
        }
      }
    }
    
    int keyboardInteractiveLogin()
    {
      int err = ssh_userauth_kbdint(socket()->sshSession(), NULL, NULL);
      char *name, *instruction, *prompt;
      int i, n;
      char echo;
      
      while (err == SSH_AUTH_INFO) {
        name = ssh_userauth_kbdint_getname(socket()->sshSession());
        instruction = ssh_userauth_kbdint_getinstruction(socket()->sshSession());
        n = ssh_userauth_kbdint_getnprompts(socket()->sshSession());
        
        // FIXME Name and instruction are currently ignored. The libssh API reference
        // suggests displaying an interactive dialog box for the user to supply the
        // information requested from the server.
            
        for(i = 0; i < n; ++i) {
          prompt = ssh_userauth_kbdint_getprompt(socket()->sshSession(), i, &echo);
          
          if (!echo) {
            // We should send the password (since only the password should be masked)
            ssh_userauth_kbdint_setanswer(socket()->sshSession(), i, (char*) socket()->getCurrentUrl().pass().ascii());
          } else {
            // FIXME Server requests something else ?
          }
        }
        
        err = ssh_userauth_kbdint(socket()->sshSession(), NULL, NULL);
      }
      
      return err;
    }
};

void SftpSocket::protoConnect(const KURL &url)
{
  emitEvent(Event::EventState, i18n("Connecting..."));
  emitEvent(Event::EventMessage, i18n("Connecting to %1:%2...").arg(url.host()).arg(url.port()));
  
  if (!getConfig("encoding").isEmpty())
    changeEncoding(getConfig("encoding"));
  
  // Connect to the remote host
  setCurrentUrl(url);
  activateCommandClass(SftpCommandConnect);
}

// *******************************************************************************************
// **************************************** DISCONNECT ***************************************
// *******************************************************************************************

void SftpSocket::protoDisconnect()
{
  Socket::protoDisconnect();
    
  if (m_sftpSession)
    sftp_free(m_sftpSession);
    
  ssh_disconnect(m_sshSession);
  m_sshSession = 0;
  
  m_login = false;
}

void SftpSocket::protoAbort()
{
  Socket::protoAbort();
  
  if (getCurrentCommand() == Commands::CmdGet || getCurrentCommand() == Commands::CmdPut) {
    // Abort current command
    resetCommandClass(UserAbort);
    emitEvent(Event::EventMessage, i18n("Aborted."));
  }
}

// *******************************************************************************************
// ******************************************* LIST ******************************************
// *******************************************************************************************

class SftpCommandList : public Commands::Base {
public:
    enum State {
      None
    };
    
    ENGINE_STANDARD_COMMAND_CONSTRUCTOR(SftpCommandList, SftpSocket, CmdList)
    
    void process()
    {
      // Check the directory listing cache
      DirectoryListing cached = Cache::self()->findCached(socket(), socket()->getCurrentDirectory());
      if (cached.isValid()) {
        socket()->emitEvent(Event::EventMessage, i18n("Using cached directory listing."));
        
        if (socket()->isChained()) {
          // We don't emit an event, because this list has been called from another
          // command. Just save the listing.
          socket()->m_lastDirectoryListing = cached;
        } else
          socket()->emitEvent(Event::EventDirectoryListing, cached);
          
        socket()->resetCommandClass();
        return;
      }
      
      socket()->m_lastDirectoryListing = DirectoryListing(socket()->getCurrentDirectory());
          
      SFTP_DIR *m_dir = sftp_opendir(socket()->sftpSession(), socket()->remoteEncoding()->encode(socket()->getCurrentDirectory()).data());
      if (!m_dir) {
        if (socket()->errorReporting()) {
          socket()->emitError(ListFailed);
          socket()->resetCommandClass(Failed);
        } else
          socket()->resetCommandClass();
        return;
      }
      
      // Read the specified directory
      SFTP_ATTRIBUTES *file;
      DirectoryEntry entry;
      
      while ((file = sftp_readdir(socket()->sftpSession(), m_dir))) {
        entry.setFilename(file->name);
        
        if (entry.filename() != "." && entry.filename() != "..") {
          entry.setFilename(socket()->remoteEncoding()->decode(entry.filename().ascii()));
          entry.setOwner(file->owner);
          entry.setGroup(file->group);
          entry.setTime(file->mtime);
          entry.setSize(file->size);
          entry.setPermissions(file->permissions);
          
          if (file->permissions & S_IFDIR)
            entry.setType('d');
          else
            entry.setType('f');
          
          socket()->m_lastDirectoryListing.addEntry(entry);
        }
        
        sftp_attributes_free(file);
      }
      
      sftp_dir_close(m_dir);
      
      // Cache the directory listing
      Cache::self()->addDirectory(socket(), socket()->m_lastDirectoryListing);
      
      if (!socket()->isChained())
        socket()->emitEvent(Event::EventDirectoryListing, socket()->m_lastDirectoryListing);
      socket()->resetCommandClass();
    }
};

void SftpSocket::protoList(const KURL &path)
{
  emitEvent(Event::EventState, i18n("Fetching directory listing..."));
  emitEvent(Event::EventMessage, i18n("Fetching directory listing..."));
  
  // Set the directory that should be listed
  setCurrentDirectory(path.path());
  
  activateCommandClass(SftpCommandList);
}

// *******************************************************************************************
// ******************************************* GET *******************************************
// *******************************************************************************************

class SftpCommandGet : public Commands::Base {
public:
    enum State {
      None,
      WaitStat,
      DestChecked
    };
    
    ENGINE_STANDARD_COMMAND_CONSTRUCTOR(SftpCommandGet, SftpSocket, CmdGet)
    
    KURL sourceFile;
    KURL destinationFile;
    filesize_t resumeOffset;
    
    void process()
    {
      switch (currentState) {
        case None: {
          // Stat source file
          resumeOffset = 0;
          sourceFile.setPath(socket()->getConfig("params.get.source"));
          destinationFile.setPath(socket()->getConfig("params.get.destination"));
          
          currentState = WaitStat;
          socket()->protoStat(sourceFile);
          break;
        }
        case WaitStat: {
          socket()->emitEvent(Event::EventState, i18n("Transfering..."));
          
          if (socket()->getStatResponse().filename().isEmpty()) {
            socket()->emitError(FileNotFound);
            socket()->resetCommandClass(Failed);
            return;
          }
          
          if (TQDir::root().exists(destinationFile.path())) {
            DirectoryListing list;
            list.addEntry(socket()->getStatResponse());
            
            currentState = DestChecked;
            socket()->emitEvent(Event::EventFileExists, list);
            return;
          } else
            TDEStandardDirs::makeDir(destinationFile.directory());
        }
        case DestChecked: {
          TQFile file;
          
          if (isWakeup()) {
            // We have been waken up because a decision has been made
            FileExistsWakeupEvent *event = static_cast<FileExistsWakeupEvent*>(m_wakeupEvent);
            
            switch (event->action) {
              case FileExistsWakeupEvent::Rename: {
                // Change the destination filename, otherwise it is the same as overwrite
                destinationFile.setPath(event->newFileName);
              }
              case FileExistsWakeupEvent::Overwrite: {
                file.setName(destinationFile.path());
                file.open(IO_WriteOnly | IO_Truncate);
                break;
              }
              case FileExistsWakeupEvent::Resume: {
                file.setName(destinationFile.path());
                file.open(IO_WriteOnly | IO_Append);
                
                // Signal resume
                resumeOffset = file.size();
                socket()->emitEvent(Event::EventResumeOffset, resumeOffset);
                break;
              }
              case FileExistsWakeupEvent::Skip: {
                // Transfer should be aborted
                socket()->emitEvent(Event::EventTransferComplete);
                socket()->resetCommandClass();
                return;
              }
            }
          } else {
            // The file doesn't exist so we are free to overwrite
            file.setName(destinationFile.path());
            file.open(IO_WriteOnly | IO_Truncate);
          }
          
          // Download the file
          SFTP_FILE *rfile = sftp_open(socket()->sftpSession(), socket()->remoteEncoding()->encode(sourceFile.path()).data(), O_RDONLY, 0);
          if (!rfile) {
            file.close();
            socket()->resetCommandClass(Failed);
            return;
          }
          
          if (resumeOffset > 0)
            sftp_seek(rfile, resumeOffset);
          
          char buffer[16384];
          int size;
          
          do {
            size = sftp_read(rfile, buffer, sizeof(buffer));
    
            if (size > 0) {
              file.writeBlock(buffer, size);
              socket()->m_transferBytes += size;
            }
            
            if (socket()->shouldAbort())
              break;
          } while (size);
          
          sftp_file_close(rfile);
          file.close();
          
          socket()->emitEvent(Event::EventTransferComplete);
          socket()->resetCommandClass();
          break;
        }
      }
    }
};

void SftpSocket::protoGet(const KURL &source, const KURL &destination)
{
  emitEvent(Event::EventState, i18n("Transfering..."));
  emitEvent(Event::EventMessage, i18n("Downloading file '%1'...").arg(source.fileName()));
  
  // Set the source and destination
  setConfig("params.get.source", source.path());
  setConfig("params.get.destination", destination.path());
  
  m_transferBytes = 0;
  
  m_speedLastTime = time(0);
  m_speedLastBytes = 0;
  
  activateCommandClass(SftpCommandGet);
}

// *******************************************************************************************
// ******************************************* PUT *******************************************
// *******************************************************************************************

class SftpCommandPut : public Commands::Base {
public:
    enum State {
      None,
      WaitStat,
      DestChecked
    };
    
    ENGINE_STANDARD_COMMAND_CONSTRUCTOR(SftpCommandPut, SftpSocket, CmdPut)
    
    KURL sourceFile;
    KURL destinationFile;
    filesize_t resumeOffset;
    
    void process()
    {
      switch (currentState) {
        case None: {
          // Stat source file
          resumeOffset = 0;
          sourceFile.setPath(socket()->getConfig("params.get.source"));
          destinationFile.setPath(socket()->getConfig("params.get.destination"));
          
          if (!TQDir::root().exists(sourceFile.path())) {
            socket()->emitError(FileNotFound);
            socket()->resetCommandClass(Failed);
            return;
          }
          
          currentState = WaitStat;
          socket()->protoStat(destinationFile);
          break;
        }
        case WaitStat: {
          socket()->emitEvent(Event::EventState, i18n("Transfering..."));
          
          if (!socket()->getStatResponse().filename().isEmpty()) {
            DirectoryListing list;
            list.addEntry(socket()->getStatResponse());
            
            currentState = DestChecked;
            socket()->emitEvent(Event::EventFileExists, list);
            return;
          } else {
            // Create destination directories
            socket()->setErrorReporting(false);
            
            TQString destinationDir = destinationFile.directory();
            TQString fullPath;
            
            for (int i = 1; i <= destinationDir.contains('/'); i++) {
              fullPath += "/" + destinationDir.section('/', i, i);
              
              // Create the directory
              socket()->protoMkdir(fullPath);
            }
          }          
        }
        case DestChecked: {
          TQFile file;
          
          if (isWakeup()) {
            // We have been waken up because a decision has been made
            FileExistsWakeupEvent *event = static_cast<FileExistsWakeupEvent*>(m_wakeupEvent);
            
            switch (event->action) {
              case FileExistsWakeupEvent::Rename: {
                // Change the destination filename, otherwise it is the same as overwrite
                destinationFile.setPath(event->newFileName);
              }
              case FileExistsWakeupEvent::Overwrite: {
                file.setName(sourceFile.path());
                file.open(IO_ReadOnly);
                break;
              }
              case FileExistsWakeupEvent::Resume: {
                resumeOffset = socket()->getStatResponse().size();
                
                file.setName(sourceFile.path());
                file.open(IO_ReadOnly);
                file.at(resumeOffset);
                
                // Signal resume
                socket()->emitEvent(Event::EventResumeOffset, resumeOffset);
                break;
              }
              case FileExistsWakeupEvent::Skip: {
                // Transfer should be aborted
                socket()->emitEvent(Event::EventTransferComplete);
                socket()->resetCommandClass();
                return;
              }
            }
          } else {
            // The file doesn't exist so we are free to overwrite
            file.setName(sourceFile.path());
            file.open(IO_ReadOnly);
          }
          
          // Download the file
          SFTP_FILE *rfile;
          
          if (resumeOffset > 0) {
            rfile = sftp_open(socket()->sftpSession(), socket()->remoteEncoding()->encode(destinationFile.path()).data(), O_WRONLY | O_APPEND, 0);
            sftp_seek(rfile, resumeOffset);
          } else
            rfile = sftp_open(socket()->sftpSession(), socket()->remoteEncoding()->encode(destinationFile.path()).data(), O_WRONLY | O_CREAT, 0);
            
          if (!rfile) {
            file.close();
            socket()->resetCommandClass(Failed);
            return;
          }
          
          char buffer[16384];
          int size;
          
          do {
            size = file.readBlock(buffer, sizeof(buffer));
    
            if (size > 0) {
              sftp_write(rfile, buffer, size);
              socket()->m_transferBytes += size;
            }
            
            if (socket()->shouldAbort())
              break;
          } while (size);
          
          sftp_file_close(rfile);
          file.close();
          
          socket()->emitEvent(Event::EventTransferComplete);
          socket()->resetCommandClass();
          break;
        }
      }
    }
};

void SftpSocket::protoPut(const KURL &source, const KURL &destination)
{
  emitEvent(Event::EventState, i18n("Transfering..."));
  emitEvent(Event::EventMessage, i18n("Uploading file '%1'...").arg(source.fileName()));
  
  // Set the source and destination
  setConfig("params.get.source", source.path());
  setConfig("params.get.destination", destination.path());
  
  m_transferBytes = 0;
  
  m_speedLastTime = time(0);
  m_speedLastBytes = 0;
  
  activateCommandClass(SftpCommandPut);
}

// *******************************************************************************************
// **************************************** REMOVE *******************************************
// *******************************************************************************************

void SftpSocket::protoRemove(const KURL &path)
{
  emitEvent(Event::EventState, i18n("Removing..."));
  
  // Remove a file or directory
  int result = 0;
  
  if (getConfigInt("params.remove.directory"))
    result = sftp_rmdir(m_sftpSession, remoteEncoding()->encode(path.path()).data());
  else
    result = sftp_rm(m_sftpSession, remoteEncoding()->encode(path.path()).data());
    
  if (result < 0) {
    resetCommandClass(Failed);
  } else {
    // Invalidate cached parent entry (if any)
    Cache::self()->invalidateEntry(this, path.directory());
    
    emitEvent(Event::EventReloadNeeded);
    resetCommandClass();
  }
}

// *******************************************************************************************
// **************************************** RENAME *******************************************
// *******************************************************************************************

void SftpSocket::protoRename(const KURL &source, const KURL &destination)
{
  emitEvent(Event::EventState, i18n("Renaming..."));
  
  if (sftp_rename(m_sftpSession, remoteEncoding()->encode(source.path()).data(), remoteEncoding()->encode(destination.path()).data()) < 0) {
    resetCommandClass(Failed);
  } else {
    // Invalidate cached parent entry (if any)
    Cache::self()->invalidateEntry(this, source.directory());
    Cache::self()->invalidateEntry(this, destination.directory());
    
    emitEvent(Event::EventReloadNeeded);
    resetCommandClass();
  }
}

// *******************************************************************************************
// **************************************** CHMOD ********************************************
// *******************************************************************************************

void SftpSocket::protoChmodSingle(const KURL &path, int mode)
{
  emitEvent(Event::EventState, i18n("Changing mode..."));
  
  SFTP_ATTRIBUTES *attrs = static_cast<SFTP_ATTRIBUTES*>(new SFTP_ATTRIBUTES);
  memset(attrs, 0, sizeof(*attrs));
  
  attrs->permissions = intToPosix(mode);
  attrs->flags = SSH_FILEXFER_ATTR_PERMISSIONS;
  
  sftp_setstat(m_sftpSession, remoteEncoding()->encode(path.path()).data(), attrs);
  sftp_attributes_free(attrs);
  
  // Invalidate cached parent entry (if any)
  Cache::self()->invalidateEntry(this, path.directory());
  
  emitEvent(Event::EventReloadNeeded);
  resetCommandClass();
}

// *******************************************************************************************
// **************************************** MKDIR ********************************************
// *******************************************************************************************

void SftpSocket::protoMkdir(const KURL &path)
{
  SFTP_ATTRIBUTES *attrs = static_cast<SFTP_ATTRIBUTES*>(new SFTP_ATTRIBUTES);
  memset(attrs, 0, sizeof(*attrs));
  
  if (sftp_mkdir(m_sftpSession, remoteEncoding()->encode(path.path()).data(), attrs) < 0) {
    if (errorReporting())
      resetCommandClass(Failed);
  } else {
    // Invalidate cached parent entry (if any)
    Cache::self()->invalidateEntry(this, path.directory());
    
    if (errorReporting()) {
      emitEvent(Event::EventReloadNeeded);
      resetCommandClass();
    }
  }
  
  delete attrs;
}

}