#    This file is part of Metasm, the Ruby assembly manipulation suite
#    Copyright (C) 2006-2009 Yoann GUILLOT
#
#    Licence is LGPL, see LICENCE in the top-level directory

require 'metasm/exe_format/main'
require 'metasm/encode'
require 'metasm/decode'

module Metasm

class JavaClass < ExeFormat
  MAGIC = "\xCA\xFE\xBA\xBE"

  CONSTANT_TAG = {0x1 => 'Utf8',   0x3 => 'Integer',
      0x4 => 'Float',  0x5 => 'Long',
      0x6 => 'Double', 0x7 => 'Class',
      0x8 => 'String', 0x9 => 'Fieldref',
      0xa => 'Methodref', 0xb => 'InterfaceMethodref',
      0xc => 'NameAndType' }

  class SerialStruct < Metasm::SerialStruct
    new_int_field :u1, :u2, :u4
  end

  class Header < SerialStruct
    mem :magic, 4, MAGIC
    u2 :minor_version
    u2 :major_version
  end

  class ConstantPool < SerialStruct
    u2 :constant_pool_count
    attr_accessor :constant_pool

    def decode(c)
      super(c)

      @constant_pool = [nil]

      i = 1
      while i < @constant_pool_count
        entry = ConstantPoolInfo.decode(c)
        entry.idx = i
        @constant_pool << entry
        i += 1

        if entry.tag =~ /Long|Double/
          # we must insert a phantom cell
          # for long and double constants
          @constant_pool << nil
          i += 1
        end
      end
    end

    def encode(c)
      cp = super(c)

      @constant_pool.each { |entry|
        next if entry.nil?
        cp << entry.encode(c)
      }
      cp
    end

    def [](idx)
      @constant_pool[idx]
    end

    def []=(idx, val)
      raise 'cannot be used to add a cp entry' if @constant_pool[idx].nil?
      @constant_pool[idx] = val
    end
  end

  class ConstantPoolInfo < SerialStruct
    u1 :tag
    fld_enum :tag, CONSTANT_TAG
    attr_accessor :info, :idx

    def decode(c)
      super(c)

      case @tag
      when 'Utf8'
        @info = ConstantUtf8.decode(c)
      when /Integer|Float/
        @info = ConstantIntFloat.decode(c)
      when /Long|Double/
        @info = ConstantLongDouble.decode(c)
      when /Class|String/
        @info = ConstantIndex.decode(c)
      when /ref$/
        @info = ConstantRef.decode(c)
      when 'NameAndType'
        @info = ConstantNameAndType.decode(c)
      else
        raise 'unknown constant tag'
        return
      end
    end

    def encode(c)
      super(c) << @info.encode(c)
    end
  end

  class ConstantUtf8 < SerialStruct
    u2 :length
    attr_accessor :bytes

    def decode(c)
      super(c)
      @bytes = c.encoded.read(@length)
    end

    def encode(c)
      super(c) << @bytes
    end
  end

  class ConstantIntFloat < SerialStruct
    u4 :bytes
  end

  class ConstantLongDouble < SerialStruct
    u4 :high_bytes
    u4 :low_bytes
  end

  class ConstantIndex < SerialStruct
    u2 :index
  end

  class ConstantRef < SerialStruct
    u2 :class_index
    u2 :name_and_type_index
  end

  class ConstantNameAndType < SerialStruct
    u2 :name_index
    u2 :descriptor_index
  end

  class ClassInfo < SerialStruct
    u2 :access_flags
    u2 :this_class
    u2 :super_class
  end

  class Interfaces < SerialStruct
    u2 :interfaces_count
    attr_accessor :interfaces

    def decode(c)
      super(c)

      @interfaces = []
      @interfaces_count.times {
        @interfaces << ConstantIndex.decode(c)
      }
    end

    def encode(c)
      ret = super(c)

      @interfaces.each { |e|
        ret << e.encode(c)
      }
      ret
    end

    def [](idx)
      @interfaces[idx]
    end
  end

  class Fields < SerialStruct
    u2 :fields_count
    attr_accessor :fields

    def decode(c)
      super(c)
      @fields = []
      @fields_count.times {
        @fields << FieldMethodInfo.decode(c)
      }
    end

    def encode(c)
      ret = super(c)

      @fields.each { |e|
        ret << e.encode(c)
      }
      ret
    end

    def [](idx)
      @fields[idx]
    end
  end

  class Methods < SerialStruct
    u2 :methods_count
    attr_accessor :methods

    def decode(c)
      super(c)
      @methods = []
      @methods_count.times {
        @methods << FieldMethodInfo.decode(c)
      }
    end

    def encode(c)
      ret = super(c)

      @methods.each { |e|
        ret << e.encode(c)
      }
      ret
    end

    def [](idx)
      @methods[idx]
    end
  end

  class FieldMethodInfo < SerialStruct
    u2 :access_flags
    u2 :name_index
    u2 :descriptor_index
    attr_accessor :attributes

    def decode(c)
      super(c)
      @attributes = Attributes.decode(c)
    end

    def encode(c)
      super(c) << @attributes.encode(c)
    end
  end

  class Attributes < SerialStruct
    u2 :attributes_count
    attr_accessor :attributes

    def decode(c)
      super(c)

      @attributes = []
      @attributes_count.times { |i|
        @attributes << AttributeInfo.decode(c)
      }
    end

    def encode(c)
      ret = super(c)

      @attributes.each { |e|
        ret << e.encode(c)
      }
      ret
    end

    def [](idx)
      @attributes[idx]
    end
  end

  class AttributeInfo < SerialStruct
    u2 :attribute_name_index
    u4 :attribute_length
    attr_accessor :data

    def decode(c)
      super(c)
      @data = c.encoded.read(@attribute_length)
    end

    def encode(c)
      super(c) << @data
    end
  end

  def encode_u1(val) Expression[val].encode(:u8, @endianness) end
  def encode_u2(val) Expression[val].encode(:u16, @endianness) end
  def encode_u4(val) Expression[val].encode(:u32, @endianness) end
  def decode_u1(edata = @encoded) edata.decode_imm(:u8, @endianness) end
  def decode_u2(edata = @encoded) edata.decode_imm(:u16, @endianness) end
  def decode_u4(edata = @encoded) edata.decode_imm(:u32, @endianness) end
  def sizeof_u1 ; 1 ; end
  def sizeof_u2 ; 2 ; end
  def sizeof_u4 ; 4 ; end

  attr_accessor :header, :constant_pool, :class_info, :interfaces, :fields, :methods, :attributes

  def initialize(endianness=:big)
    @endianness = endianness
    @encoded = EncodedData.new
    super()
  end

  def decode
    @header = Header.decode(self)
    @constant_pool = ConstantPool.decode(self)
    @class_info = ClassInfo.decode(self)
    @interfaces = Interfaces.decode(self)
    @fields = Fields.decode(self)
    @methods = Methods.decode(self)
    @attributes = Attributes.decode(self)
  end

  def encode
    @encoded = EncodedData.new
    @encoded << @header.encode(self)
    @encoded << @constant_pool.encode(self)
    @encoded << @class_info.encode(self)
    @encoded << @interfaces.encode(self)
    @encoded << @fields.encode(self)
    @encoded << @methods.encode(self)
    @encoded << @attributes.encode(self)
    @encoded.data
  end

  def cpu_from_headers
    raise 'JVM'
  end

  def each_section
    raise 'n/a'
  end

  def get_default_entrypoints
    []
  end

  def string_at(idx)
    loop do
      tmp = @constant_pool[idx].info
      return tmp.bytes if tmp.kind_of? ConstantUtf8
      idx = tmp.index
    end
  end

  def decode_methodref(mref)
    class_idx = mref.info.class_index
    nt_idx = mref.info.name_and_type_index
    name_idx = @constant_pool[nt_idx].info.name_index
    desc_idx = @constant_pool[nt_idx].info.descriptor_index

    string_at(class_idx) + '/' + string_at(name_idx) + string_at(desc_idx)
  end

  def cp_add(cpi, tag)
    cpe = ConstantPoolInfo.new
    cpe.tag = tag
    cpe.info = cpi
    cpe.idx = @constant_pool.constant_pool_count

    @constant_pool.constant_pool << cpe
    @constant_pool.constant_pool_count += 1
    @constant_pool.constant_pool_count += 1 if tag =~ /Long|Double/

    cpe.idx
  end

  def cp_find(tag)
    constant_pool.constant_pool.each { |e|
      next if !e or e.tag != tag
      if yield(e.info)
        return e.idx
      end
    }
    nil
  end


  def cp_auto_utf8(string)
    if idx = cp_find('Utf8') { |i| i.bytes == string }
      return idx
    end

    cpi = ConstantUtf8.new
    cpi.bytes = string
    cpi.length = string.length
    cp_add(cpi, 'Utf8')
  end

  def cp_auto_class(classname)
    if idx = cp_find('Class') { |i| string_at(i.index) == classname }
      return idx
    end

    cpi = ConstantIndex.new
    cpi.index = cp_auto_utf8(classname)
    cp_add(cpi, 'Class')
  end

  def cp_add_methodref(classname, name, descriptor)
    nat = ConstantNameAndType.new
    nat.name_index = cp_auto_utf8(name)
    nat.descriptor_index = cp_auto_utf8(descriptor)
    natidx = cp_add(nat, 'NameAndType')

    cpi = ConstantRef.new
    cpi.class_index = cp_auto_class(classname)
    cpi.name_and_type_index = natidx

    cp_add(cpi, 'Methodref')
  end

  def attribute_create(name, data)
    a = AttributeInfo.new
    a.attribute_name_index = cp_auto_utf8(name)
    a.attribute_length = data.size
    a.data = data
    a
  end
end
end
