require "cgi"
require "net/http"
require "net/ssh"
require "net/sftp"
require "uri"
require "uri/sftp"
require "digest/md5"
require "digest/sha1"
require "facet/progressbar"
require "highline"


# Monkeypatching: SFTP never defines the mkdir method on its session or the underlying
# driver, it just redirect calls through method_missing. Rake, on the other hand, decides
# to define mkdir on Object, and so routes our calls to FileUtils.
module Net #:nodoc:all
  class Session
    def mkdir(path, attrs = {})
      method_missing :mkdir, path, attrs
    end
  end

  class SFTP::Protocol::Driver
    def mkdir(first, path, attrs = {})
      method_missing :mkdir, first, path, attrs
    end
  end
end


module Buildr

  # Transports are used for downloading artifacts from remote repositories, uploading
  # artifacts to deployment repositories, and anything else you need to move around.
  #
  # The HTTP transport is used for all URLs with the scheme http or https. You can only
  # use the HTTP transport to download artifacts.
  #
  # The SFTP transport is used for all URLs with the schema sftp. You can only use the
  # SFTP transport to upload artifacts.
  #
  # The SFTP transport supports the following options:
  # * :username -- The username.
  # * :password -- A password. If unspecified, you will be prompted to enter a password.
  # * :permissions -- Permissions to set on the uploaded file.
  # You can also pass the username/password in the URL.
  #
  # The SFTP transport will automatically create MD5 and SHA1 digest files for each file
  # it uploads.
  module Transports

    # Indicates the requested resource was not found.
    class NotFound < Exception
    end

    class << self

      # :call-seq:
      #   perform(url, options?) { |transport| ... }
      #
      # Perform one or more operations using an open connection to the
      # specified URL. For examples, see Transport#download and Transport#upload.
      def perform(url, options = nil, &block)
        uri = URI.parse(url.to_s)
        const_get(uri.scheme.upcase).perform(uri, options, &block)
      end

      # :call-seq:
      #   download(url, target, options?)
      #
      # Convenience method for downloading a single file from the specified
      # URL to the target file.
      def download(url, target, options = nil)
        uri = URI.parse(url.to_s)
        path, uri.path = uri.path, ""
        const_get(uri.scheme.upcase).perform(uri, options) do |transport|
          transport.download(path, target)
        end
      end

    end

    # Extend this class if you are implementing a new transport.
    class Transport

      class << self

        # :call-seq:
        #   perform(url, options?) { |transport| ... }
        #
        # Perform one or more operations using an open connection to the
        # specified URL. For examples, see #download and #upload.
        def perform(url, options = nil)
          instance = new(url, options)
          begin
            yield instance
          ensure
            instance.close
          end
        end
      end

      # The server URI.
      attr_reader :uri
      # The base path on the server, always ending with a slash.
      attr_reader :base_path
      # Options passed during construction.
      attr_reader :options

      # Initialize the transport with the specified URL and options.
      def initialize(url, options)
        @uri = URI.parse(url.to_s)
        @base_path = @uri.path || "/"
        @base_path += "/" unless @base_path[-1] == ?/
        @options = options || {}
      end

      # Downloads a file from the specified path, relative to the server URI.
      # Downloads to either the target file, or by calling the block with each
      # chunk of the file. Returns the file's modified timestamp, or now.
      #
      # For example:
      #   Transports.perform("http://server/libs") do |http|
      #     http.download("my_project/test.jar", "test.jar")
      #     http.download("my_project/readme") { |text| $stdout.write text }
      #   end
      def download(path, target, &block)
        fail "Upload not implemented for this transport"
      end

      # Uploads a file (source) to the server at the specified path,
      # relative to the server URI.
      #
      # For example:
      #   Transports.perform("sftp://server/libs") do |sftp|
      #     sftp.mkpath "my_project"
      #     sftp.upload("target/test.jar", "my_project/test.jar")
      #   end
      def upload(source, path)
        fail "Upload not implemented for this transport"
      end

      # Creates a path on the server relative to the server URI.
      # See #upload for example.
      def mkpath(path)
        fail "Upload not implemented for this transport"
      end

    protected

      # :call-seq:
      #   with_progress_bar(file_name, length) { |progress| ... }
      #
      # Displays a progress bar while executing the block. The first
      # argument provides a filename to display, the second argument
      # its size in bytes.
      #
      # The block is yielded with a progress object that implements
      # a single method. Call << for each block of bytes down/uploaded.
      def with_progress_bar(file_name, length)
        if verbose && $stdout.isatty
          progress_bar = Console::ProgressBar.new(file_name, length)
          # Extend the progress bar so we can display count/total.
          class << progress_bar
            def total()
              convert_bytes(@total)
            end
          end
          # Squeeze the filename into 30 characters.
          if file_name.size > 30
            base, ext = file_name.split(".")
            truncated = "#{base[0..26-ext.size]}...#{ext}"
          else
            truncated = file_name
          end
          progress_bar.format = "#{truncated}: %3d%% %s %s/%s %s"
          progress_bar.format = "%3d%% %s %s/%s %s"
          progress_bar.format_arguments = [:percentage, :bar, :bytes, :total, :stat]
          progress_bar.bar_mark = "."


          begin
            class << progress_bar
              def <<(bytes)
                inc bytes.respond_to?(:size) ? bytes.size : bytes
              end
            end
            yield progress_bar
          ensure
            progress_bar.finish
          end
        else
          progress_bar = Object.new
          class << progress_bar
            def <<(bytes)
            end
          end
          yield progress_bar
        end
      end

      # :call-seq:
      #   with_digests(types?) { |digester| ... } => hash
      #
      # Use the Digester to create digests for files you are downloading or
      # uploading, and either verify their signatures (download) or create
      # signatures (upload).
      #
      # The method takes one argument with the list of digest algorithms to
      # support. Leave if empty and it will default to MD5 and SHA1.
      #
      # The method then yields the block passing it a Digester. The Digester
      # supports two methods. Use << to pass data that you are down/uploading.
      # Once all data is transmitted, use each to iterate over the digests.
      # The each method calls the block with the digest type (e.g. "md5")
      # and the hexadecimal digest value.
      #
      # For example:
      #   with_digests do |digester|
      #     download url do |block|
      #       digester << block
      #     end
      #     digester.each do |type, hexdigest|
      #       signature = download "#{url}.#{type}"
      #       fail "Mismatch" unless signature == hexdigest
      #     end
      #   end
      def with_digests(types = nil)
        digester = Digester.new(types)
        yield digester
        digester.to_hash
      end

      class Digester #:nodoc:

        def initialize(types)
          types ||= [ "md5", "sha1" ]
          @digests = types.inject({}) do |hash, type|
            hash[type.to_s.downcase] = Digest.const_get(type.to_s.upcase).new
            hash
          end
        end

        # Add bytes for digestion.
        def <<(bytes)
          @digests.each { |type, digest| digest << bytes }
        end

        # Iterate over all the digests calling the block with two arguments:
        # the digest type (e.g. "md5") and the hexadecimal digest value.
        def each()
          @digests.each { |type, digest| yield type, digest.hexdigest }
        end

        # Returns a hash that maps each digest type to its hexadecimal digest value.
        def to_hash()
          @digests.keys.inject({}) do |hash, type|
            hash[type] = @digests[type].hexdigest
            hash
          end
        end

      end

    end


    class HTTP < Transport #:nodoc:

      def initialize(url, options)
        super
        @http = Net::HTTP.start(@uri.host, @uri.port)
      end

      def download(path, target = nil, &block)
        puts "Requesting #{@uri}/#{path} " if Rake.application.options.trace
        last_modified = File.stat(target).mtime.utc if target && File.exist?(target)
        headers = {}
        headers["If-Modified-Since"] = CGI.rfc1123_date(last_modified) if last_modified
        @http.request_get(@base_path + path, headers) do |response|
          case response
          when Net::HTTPNotModified
            # No modification, nothing to do.
            puts "Not modified since last download" if Rake.application.options.trace

          when Net::HTTPRedirection
            # Try to download from the new URI, handle relative redirects.
            puts "Redirected to #{response['Location']}" if Rake.application.options.trace
            last_modified = Transports.download(@uri + URI.parse(response["location"]), target, @options)

          when Net::HTTPOK
            puts "Downloading #{@uri}/#{path}" if verbose
            last_modified = Time.parse(response["Last-Modified"] || "")
            with_progress_bar path.split("/").last, response.content_length do |progress|
              with_digests(@options[:digests]) do |digester|

                download = proc do |write|
                  # Read the body of the page and write it out.
                  response.read_body do |chunk|
                    write[chunk]
                    digester << chunk
                    progress << chunk
                  end
                  # Check server digests before approving the download.
                  digester.each do |type, hexdigest|
                    @http.request_get("#{@base_path}#{path}.#{type.to_s.downcase}") do |response|
                      if Net::HTTPOK === response
                        puts "Checking signature from #{@uri}/#{path}" if Rake.application.options.trace
                        fail "Checksum failure for #{@uri}/#{path}: #{type.to_s.upcase} digest on server did not match downloaded file" unless
                          response.read_body.split.first == hexdigest
                      end
                    end
                  end
                end

                if target
                  # If download breaks we end up with a partial file which is
                  # worse than not having a file at all, so download to temporary
                  # file and then move over.
                  temp = Tempfile.new(File.basename(target))
                  download[ proc { |chunk| temp.write chunk } ]
                  temp.close
                  File.move temp.path, target
                  File.utime last_modified, last_modified, target
                else
                  download[ block ]
                end

              end
            end
          when Net::HTTPNotFound
            raise NotFound
          else
            fail "Failed to download #{@uri}/#{path}: #{response.message}"
          end
        end
        last_modified
      end

      def close()
        @http.finish
        @http = nil
      end

    end

    # Use the HTTP transport for HTTPS connections.
    HTTPS = HTTP #:nodoc:


    class SFTP < Transport #:nodoc:

      class << self
        def passwords()
          @passwords ||= {}
        end
      end

      attr_reader :sftp

      def initialize(url, options)
        super
        @permissions = options.delete :permissions
        # SSH options are based on the username/password from the URI.
        ssh_options = { :port=>@uri.port, :username=>@uri.user }.merge(options || {})
        ssh_options[:password] ||= SFTP.passwords[@uri.host]
        begin
          puts "Connecting to #{@uri.host}" if Rake.application.options.trace
          session = Net::SSH.start(@uri.host, ssh_options)
          SFTP.passwords[@uri.host] = ssh_options[:password]
        rescue Net::SSH::AuthenticationFailed=>ex
          # Only if running with console, prompt for password.
          if !ssh_options[:password] && $stdout.isatty
            password = HighLine.new.ask("Password for #{@uri.host}:") { |q| q.echo = "*" }
            ssh_options[:password] = password
            retry
          end
          raise
        end
        @sftp = session.sftp.connect
        puts "connected" if Rake.application.options.trace
      end

      def upload(source, path)
        File.open(source) do |file|
          with_progress_bar path.split("/").last, File.size(source) do |progress|
            with_digests(@options[:digests]) do |digester|
              target_path = "#{@base_path}#{path}"
              puts "Uploading to #{target_path}" if Rake.application.options.trace
              @sftp.open_handle(target_path, "w") do |handle|
                # Writing in chunks gives us the benefit of a progress bar,
                # but also require that we maintain a position in the file,
                # since write() with two arguments always writes at position 0.
                pos = 0
                while chunk = file.read(32 * 4096)
                  @sftp.write(handle, chunk, pos)
                  pos += chunk.size
                  digester << chunk
                  progress << chunk
                end
              end
              @sftp.setstat(target_path, :permissions => @permissions) if @permissions

              # Upload all the digests.
              digester.each do |type, hexdigest|
                digest_file = "#{@base_path}#{path}.#{type}"
                puts "Uploading signature to #{digest_file}" if Rake.application.options.trace
                @sftp.open_handle(digest_file, "w") do |handle|
                  @sftp.write(handle, "#{hexdigest}  #{path}")
                end
                @sftp.setstat(digest_file, :permissions => @permissions) if @permissions
              end
            end

          end
        end
      end

      def mkpath(path)
        # To create a path, we need to create all its parent.
        # We use realpath to determine if the path already exists,
        # otherwise mkdir fails.
        puts "Creating path #{@base_path}" if Rake.application.options.trace
        path.split("/").inject(@base_path) do |base, part|
          combined = base + part
          @sftp.realpath combined rescue @sftp.mkdir combined, {}
          "#{combined}/"
        end
      end

      def close()
        @sftp.close
        @sftp = nil
      end

    end

  end
end