#    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/main'
require 'metasm/parse'
require 'metasm/encode'
require 'metasm/decode'
require 'metasm/exe_format/serialstruct'
require 'metasm/os/main'	# VirtualFile

module Metasm
class ExeFormat
  # creates a new instance, populates self.encoded with the supplied string
  def self.load(str, *a, &b)
    e = new(*a, &b)
    if str.kind_of? EncodedData; e.encoded = str
    else e.encoded << str
    end
    e
  end

  # same as load, used by AutoExe
  def self.autoexe_load(*x, &b)
    load(*x, &b)
  end

  attr_accessor :filename

  # same as +load+, but from a file
  # uses VirtualFile if available
  def self.load_file(path, *a, &b)
    e = load(VirtualFile.read(path), *a, &b)
    e.filename ||= path
    e
  end

  # +load_file+ then decode
  def self.decode_file(path, *a, &b)
    e = load_file(path, *a, &b)
    e.decode if not e.instance_variables.map { |iv| iv.to_s }.include?("@disassembler")
    e
  end

  # +load_file+ then decode header
  def self.decode_file_header(path, *a, &b)
    e = load_file(path, *a, &b)
    e.decode_header
    e
  end

  def self.decode(raw, *a, &b)
    e = load(raw, *a, &b)
    e.decode
    e
  end

  def self.decode_header(raw, *a, &b)
    e = load(raw, *a, &b)
    e.decode_header
    e
  end

  # creates a new object using the specified cpu, parses the asm source, and assemble
  def self.assemble(cpu, source, file='<unk>', lineno=1)
    source, cpu = cpu, source if source.kind_of? CPU
    e = new(cpu)
    e.assemble(source, file, lineno)
    e
  end

  # same as #assemble, reads asm source from the specified file
  def self.assemble_file(cpu, filename)
    filename, cpu = cpu, filename if filename.kind_of? CPU
    assemble(cpu, File.read(filename), filename, 1)
  end

  # parses a bunch of standalone C code, compile and assemble it
  def compile_c(source, file='<unk>', lineno=1)
    cp = @cpu.new_cparser
    tune_cparser(cp)
    cp.parse(source, file, lineno)
    read_c_attrs cp if respond_to? :read_c_attrs
    asm_source = @cpu.new_ccompiler(cp, self).compile
    puts asm_source if $DEBUG
    assemble(asm_source, 'C compiler output', 1)
    c_set_default_entrypoint
  end

  # creates a new object using the specified cpu, parse/compile/assemble the C source
  def self.compile_c(cpu, source, file='<unk>', lineno=1)
    source, cpu = cpu, source if source.kind_of? CPU
    e = new(cpu)
    e.compile_c(source, file, lineno)
    e
  end

  def self.compile_c_file(cpu, filename)
    filename, cpu = cpu, filename if filename.kind_of? CPU
    compile_c(cpu, File.read(filename), filename, 1)
  end

  # add directive to change the current assembler section to the assembler source +src+
  def compile_setsection(src, section)
    src << section
  end

  # prepare a preprocessor before it reads any source, should define macros to identify the fileformat
  def tune_prepro(l)
  end

  # prepare a cparser
  def tune_cparser(cp)
    tune_prepro(cp.lexer)
  end

  # this is called once C code is parsed, to handle C attributes like export/import/init etc
  def read_c_attrs(cp)
  end

  # should setup a default entrypoint for C code, including preparing args for main() etc
  def c_set_default_entrypoint
  end

  attr_writer :disassembler	# custom reader
  def disassembler
    @disassembler ||= init_disassembler
  end

  # returns the exe disassembler
  # if it does not exist, creates one, and feeds it with the exe sections
  def init_disassembler
    @disassembler ||= Disassembler.new(self)
    @disassembler.cpu ||= cpu
    each_section { |edata, base|
      edata ||= EncodedData.new
      @disassembler.add_section edata, base
    }
    @disassembler
  end

  # disassembles the specified entrypoints
  # initializes the disassembler if needed
  # uses get_default_entrypoints if the argument list is empty
  # returns the disassembler
  def disassemble(*entrypoints)
    entrypoints = get_default_entrypoints if entrypoints.empty?
    disassembler.disassemble(*entrypoints)
    @disassembler
  end

  # disassembles the specified entrypoints without backtracking
  # initializes the disassembler if needed
  # uses get_default_entrypoints if the argument list is empty
  # returns the disassembler
  def disassemble_fast_deep(*entrypoints)
    entrypoints = get_default_entrypoints if entrypoints.empty?
    disassembler.disassemble_fast_deep(*entrypoints)
    @disassembler
  end

  # returns a list of entrypoints to disassemble (program entrypoint, exported functions...)
  def get_default_entrypoints
    []
  end

  # encodes the executable as a string, checks that all relocations are
  # resolved, and returns the raw string version
  def encode_string(*a)
    encode(*a)
    raise ["Unresolved relocations:", @encoded.reloc.map { |o, r| "#{r.target} " + (Backtrace.backtrace_str(r.backtrace) if r.backtrace).to_s }].join("\n") if not @encoded.reloc.empty?
    @encoded.data
  end

  # saves the result of +encode_string+ in the specified file
  # fails if the file already exists
  def encode_file(path, *a)
    #raise Errno::EEXIST, path if File.exist? path	# race, but cannot use O_EXCL, as O_BINARY is not defined in ruby
    encode_string(*a)
    File.open(path, 'wb') { |fd| fd.write(@encoded.data) }
  end

  # returns the address at which a given file offset would be mapped
  def addr_to_fileoff(addr)
    addr
  end

  # returns the file offset where a mapped byte comes from
  def fileoff_to_addr(foff)
    foff
  end

  def shortname; self.class.name.split('::').last.downcase; end

module IntToHash
  # converts a constant name to its numeric value using the hash
  # {1 => 'toto', 2 => 'tata'}: 'toto' => 1, 42 => 42, 'tutu' => raise
  def int_from_hash(val, hash)
    val.kind_of?(Integer) ? hash.index(val) || val : hash.index(val) or raise "unknown constant #{val.inspect}"
  end

  # converts an array of flag constants to its numeric value using the hash
  # {1 => 'toto', 2 => 'tata'}: ['toto', 'tata'] => 3, 'toto' => 2, 42 => 42
  def bits_from_hash(val, hash)
    val.kind_of?(Array) ? val.inject(0) { |val_, bitname| val_ | int_from_hash(bitname, hash) } : int_from_hash(val, hash)
  end

  # converts a numeric value to the corresponding constant name using the hash
  # {1 => 'toto', 2 => 'tata'}: 1 => 'toto', 42 => 42, 'tata' => 'tata', 'tutu' => raise
  def int_to_hash(val, hash)
    val.kind_of?(Integer) ? hash.fetch(val, val) : (hash.index(val) ? val : raise("unknown constant #{val.inspect}"))
  end

  # converts a numeric value to the corresponding array of constant flag names using the hash
  # {1 => 'toto', 2 => 'tata'}: 5 => ['toto', 4]
  def bits_to_hash(val, hash)
    (val.kind_of?(Integer) ? (hash.find_all { |k, v| val & k == k and val &= ~k }.map { |k, v| v } << val) : val.kind_of?(Array) ? val.map { |e| int_to_hash(e, hash) } : [int_to_hash(val, hash)]) - [0]
  end
end
  include IntToHash
end

class SerialStruct
  include ExeFormat::IntToHash
end
end
