# Copyright 2008-2009 Amazon.com, Inc. or its affiliates.  All Rights
# Reserved.  Licensed under the Amazon Software License (the
# "License").  You may not use this file except in compliance with the
# License. A copy of the License is located at
# http://aws.amazon.com/asl or in the "license" file accompanying this
# file.  This file is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See
# the License for the specific language governing permissions and
# limitations under the License.

require 'ec2/amitools/crypto'
require 'ec2/amitools/format'
require 'ec2/amitools/xmlutil'
require 'pathname'
require 'rexml/document'

# Manifest Version 2007-08-29.
# Not backwards compatible
class ManifestV20070829
  VERSION_STRING = '2007-08-29'
  VERSION = VERSION_STRING.gsub('-','_').to_i

  # Expose the version
  def self.version
    VERSION
  end
  
  # AMI part information container.
  class PartInformation
    attr_reader :filename # Part's file basename.
    attr_reader :digest   # Part's digest hex encoded.
    
    # Initialize with part's _filename_ and _digest_ as byte string.
    def initialize(filename, digest)
      @filename , @digest = filename, digest
    end
  end

  def initialize(xml = nil)
    if xml == nil
      @doc = REXML::Document.new
    else
      # Convert to string if necessary.
      xml = (xml.kind_of?(IO) ? xml.read : xml)
      @doc = REXML::Document.new(xml)
    end
  end

  # for debugging only
  def doc
    @doc
  end
  
  # Initialize the manifest with AMI information.
  # Return +true+ if the initialization was succesful.
  # Raise an exception on error.
  def init(
    name,
    user,                   # The user's account number.
    arch,                   # Target architecture for AMI.
    reserved,               # Reserved for future use; pass nil.
    parts,                  # A list of parts filenames and digest pairs.
    size,                   # The size of the AMI in bytes.
    bundled_size,           # The size of the bunled AMI in bytes.
    user_encrypted_key,     # Hex encoded.
    ec2_encrypted_key,      # Hex encoded.
    cipher_algorithm,       # The cipher algorithm used to encrypted the AMI.
    user_encrypted_iv,      # Hex encoded.
    ec2_encrypted_iv,       # Hex encoded.
    digest,                 # Hex encoded.
    digest_algorithm,       # The digest algorithm.
    privkey_filename,       # The user's private key filename.
    bundler_name = nil,
    bundler_version = nil,
    bundler_release = nil)

    # Check reserved parameters are nil
    raise ArgumentError.new("reserved parameters not nil") unless reserved.nil?

    # Check non-String parameter types.
    raise ArgumentError.new("parts parameter type invalid") unless parts.is_a? Array
        
    # XML document.
    @doc = REXML::Document.new
    @doc << REXML::XMLDecl.new

    # manifest - the root element.
    manifest = REXML::Element.new('manifest')

    @doc.add_element(manifest)
    
    # version - indicate the manifest version.
    version = REXML::Element.new('version')
    version.text = VERSION
    manifest.add_element(version)
    
    # bundler information
    if bundler_name or bundler_version or bundler_release
      bundler_element = REXML::Element.new('bundler')
      manifest.add_element(bundler_element)
      
      [['name', bundler_name ], 
       ['version', bundler_version ],
       ['release', bundler_release ]].each do |element_name, text|
        if element_name
          element = REXML::Element.new(element_name)
          element.text = text
          bundler_element.add_element(element)
        end
      end 
    end
    
    # machine_configuration - the target hardware description of the AMI.
    machine_configuration = REXML::Element.new('machine_configuration')
    manifest.add_element(machine_configuration)

    # image - the image element.
    image = REXML::Element.new('image')
    name_element = REXML::Element.new('name')
    name_element.text = name
    image.add_element(name_element)
    manifest.add_element(image)
    
    # user - the user's AWS access key ID.
    user_element = REXML::Element.new('user')
    user_element.text = user
    image.add_element(user_element)

    # m_c arch - the target hardware architecture
    if arch
      arch_element = REXML::Element.new('architecture')
      arch_element.text = arch
      machine_configuration.add_element(arch_element)
    end

    # digest - the digest of the AMI.
    digest_element = REXML::Element.new('digest')
    digest_element.add_attribute('algorithm', digest_algorithm)
    digest_element.add_text(digest)
    image.add_element(digest_element)
    
    # size - the size of the uncompressed AMI.
    size_element = REXML::Element.new('size')
    size_element.text = size.to_s
    image.add_element(size_element)
    
    # bundled size - the size of the bundled AMI.
    bundled_size_element = REXML::Element.new('bundled_size')
    bundled_size_element.text = bundled_size.to_s
    image.add_element(bundled_size_element)
    
    # ec2 encrypted key element.
    ec2_encrypted_key_element = REXML::Element.new('ec2_encrypted_key')
    ec2_encrypted_key_element.add_attribute('algorithm', cipher_algorithm)
    ec2_encrypted_key_element.add_text(ec2_encrypted_key)
    image.add_element(ec2_encrypted_key_element)
    
    # user encrypted key element.
    user_encrypted_key_element = REXML::Element.new('user_encrypted_key')
    user_encrypted_key_element.add_attribute('algorithm', cipher_algorithm)
    user_encrypted_key_element.add_text(user_encrypted_key)
    image.add_element(user_encrypted_key_element)
    
    # ec2 encrypted iv element.
    ec2_encrypted_iv_element = REXML::Element.new('ec2_encrypted_iv')
    ec2_encrypted_iv_element.add_text(ec2_encrypted_iv)
    image.add_element(ec2_encrypted_iv_element)
    
    # user encrypted iv element.
    user_encrypted_iv_element = REXML::Element.new('user_encrypted_iv')
    user_encrypted_iv_element.add_text(user_encrypted_iv)
    image.add_element(user_encrypted_iv_element)
    
    # parts - list of the image parts.
    parts_element = REXML::Element.new('parts')
    parts_element.add_attributes({'count' => parts.size.to_s})
    index=0
    parts.each do |part|
      # Add image part element for each image part.
      part_element = REXML::Element.new('part')
      part_element.add_attribute('index', index.to_s)
      filename = REXML::Element.new('filename')
      filename.add_text(part[0])
      part_element.add_element(filename)
      digest = REXML::Element.new('digest')
      digest.add_attribute('algorithm', digest_algorithm)
      digest.add_text(Format::bin2hex(part[1]))
      part_element.add_element(digest)
      parts_element.add_element(part_element)
      index+=1
    end
    image.add_element(parts_element)

    # Sign the manifest.
    sign(privkey_filename)
    
    return true
  end
  
  def self::version20070829?(xml)
    doc = REXML::Document.new(xml)
    version = REXML::XPath.first(doc.root, 'version')
    return true if (version and version.text and version.text == VERSION_STRING)
    return (version and version.text and version.text == VERSION.to_s)
  end
  
  # Return the AMI's digest hex encoded.
  def digest()
    return get_element_text('image/digest')
  end
    
  def name()
    return get_element_text('image/name')
  end

  # The ec2 encrypted key hex encoded.
  def ec2_encrypted_key()
    return get_element_text('image/ec2_encrypted_key')
  end

  # The user encrypted key hex encoded.
  def user_encrypted_key()
    return get_element_text('image/user_encrypted_key')
  end

  # The ec2 encrypted initialization vector hex encoded.
  def ec2_encrypted_iv()
    return get_element_text('image/ec2_encrypted_iv')
  end

  # The user encrypted initialization vector hex encoded.
  def user_encrypted_iv()
    return get_element_text('image/user_encrypted_iv')
  end

  # Get digest algorithm used.
  def digest_algorithm()
    return REXML::XPath.first(@doc.root, 'image/digest/@algorithm').to_s
  end

  # Get cipher algorithm used.
  def cipher_algorithm()
    return REXML::XPath.first(@doc.root, 'image/ec2_encrypted_key/@algorithm').to_s
  end

  # Retrieve a list of AMI bundle parts info. Each element is a hash
  # with the following elements:
  # * 'digest'
  # * 'filename'
  # * 'index'
  def ami_part_info_list
    parts = Array.new
    REXML::XPath.each(@doc.root,'image/parts/part') do |part|
      index = part.attribute('index').to_s.to_i
      filename = REXML::XPath.first(part, 'filename').text
      digest = REXML::XPath.first(part, 'digest').text
      parts << { 'digest'=>digest, 'filename'=>filename, 'index'=>index }
    end
    return parts
  end

  # A list of PartInformation instances representing the AMI parts.
  def parts()
    parts = []
    REXML::XPath.each(@doc.root,'image/parts/part') do |part|
      index = part.attribute('index').to_s.to_i
      filename = REXML::XPath.first(part, 'filename').text
      digest = Format::hex2bin(REXML::XPath.first(part, 'digest').text)
      parts[index] = PartInformation.new(filename, digest)
    end
    return parts
  end

  # Return the size of the AMI.
  def size()
    return get_element_text('image/size').to_i()
  end
  
  # Return the (optional) architecture of the AMI.
  def arch()
    return get_element_text('machine_configuration/architecture')
  end
  
  # Return the bundled size of the AMI.
  def bundled_size()
    return get_element_text('image/bundled_size').to_i
  end
  
  # Return the bundler name.
  def bundler_name()
    return get_element_text('bundler/name')
  end
  
  # Return the bundler version.
  def bundler_version()
    return get_element_text('bundler/version')
  end
  
  # Return the bundler release.
  def bundler_release()
    return get_element_text('bundler/release')
  end
  
  # Sign the manifest. If it is already signed, the signature and certificate
  # will be replaced
  def sign(privkey_filename)
    unless privkey_filename.kind_of? String and File::exist?(privkey_filename)
      raise ArgumentError.new("privkey_filename parameter invalid")
    end
    
    # Get the XML for <machine_configuration> and <image> elements and sign them.
    machine_configuration_xml = XMLUtil.get_xml(@doc.to_s, 'machine_configuration') || ""
    image_xml = XMLUtil.get_xml(@doc.to_s, 'image')
    sig = Crypto::sign(machine_configuration_xml + image_xml, privkey_filename)
         
    # Create the signature and certificate elements.
    signature = REXML::Element.new('signature')
    signature.add_text(Format::bin2hex(sig))
    @doc.root.delete_element('signature')
    @doc.root.add_element(signature)
  end

  # Return the signature
  def signature
    get_element_text('signature')
  end

  # Verify the signature
  def authenticate(cert)
    machine_configuration_xml = XMLUtil.get_xml(@doc.to_s, 'machine_configuration') || ""
    image_xml = XMLUtil.get_xml(@doc.to_s, 'image')
    pubkey = Crypto::cert2pubkey(cert)
    Crypto::authenticate(machine_configuration_xml + image_xml, Format::hex2bin(signature), pubkey)
  end

  # Return the manifest as an XML string.
  def to_s()
    return @doc.to_s
  end

  def user()
    return get_element_text('image/user')
  end

  def version()
    return get_element_text('version').to_i
  end
  
  private
  
  def get_element_text(xpath)
    element = REXML::XPath.first(@doc.root, xpath)
    unless element
      raise "invalid AMI manifest, #{xpath} element not present"
    end
    unless element.text
      raise "invalid AMI manifest, #{xpath} element empty"
    end
    return element.text
  end
end
