compiler.rb

# -*- coding: utf-8 -*-

require 'yaml'
require 'pp'

# Object の拡張
class Object
  # [ xxx ].pack を xxx.pack と書けるようにする
  def pack(*args)
    [ self ].pack *args
  end
end

# String の拡張
class String
  # offset 文字目から val の長さぶん val で置き換える
  def write(offset, val)
    self[offset, val.length] = val
  end

  # PE 出力用のバイト列
  # 文字列はそのままの値を使用する。
  def raw
    if self.respond_to?(:force_encoding)
      # 1.9.x では UTF-8 と ASCII の結合ができないのでその対策
      self.force_encoding("ASCII-8bit")
    else
      self
    end
  end
end

# Array の拡張
class Array
  # PE 出力用のバイト列
  # 配列は各要素の raw を単純に連結したものとする
  def raw
    # 後で値を追加することもあるので、値はキャッシュしない
    self.map{|item| item.raw }.join("")
  end

  # pack の各要素が構造体の場合は raw に変換＆バイト単位に分割してから pack するよう処理を変更
  alias :real_pack :pack
  def pack(*args)
    self.map{|elem| elem.class.ancestors.include?(Struct) ? elem.raw.bytes.to_a : elem }.flatten.real_pack(*args)
  end
end

# String の拡張
class Struct
  class << self
    # 構造体ごとにフィールドの定義を保持するためのクラスインスタンス変数
    attr_accessor :fields
  end

  # 一度計算していた場合はキャッシュから取り出す
  def raw
    @raw ||= get_raw
  end

  # PE 出力用のバイト列
  # fields の定義に従い各フィールドの pack を行い結合したもの
  def get_raw
    ret = ""
    self.class.fields.each do |field|
      # デフォルト値が Proc の場合はその処理を実行した戻り値を利用する
      # （structs の定義の段階では別の構造体が定義されていないため）
      default = field[:default]
      default = default.call if default.class == Proc

      # 一度デフォルト値の pack を行い、値が指定されていた場合はその値を pack したもので先頭から置き換える
      # （10 要素の配列中 .yml で 2 要素しか指定されていないような場合も pack できるようにするため）
      raw = default.pack(field[:template])
      if self[field[:name]]
        raw.write 0, self[field[:name]].pack(field[:template])
      end
      ret << raw
    end
    ret
  end
end

# 必要な構造体の情報を配列で定義しておく
structs = {
  "ImageDosHeader" => [
    { :name => :e_magic,    :template => "a2", :default => "" },
    { :name => :e_cblp,     :template => "S",  :default => 0 },
    { :name => :e_cp,       :template => "S",  :default => 0 },
    { :name => :e_crlc,     :template => "S",  :default => 0 },
    { :name => :e_cparhdr,  :template => "S",  :default => 0 },
    { :name => :e_minalloc, :template => "S",  :default => 0 },
    { :name => :e_maxalloc, :template => "S",  :default => 0 },
    { :name => :e_ss,       :template => "S",  :default => 0 },
    { :name => :e_sp,       :template => "S",  :default => 0 },
    { :name => :e_csum,     :template => "S",  :default => 0 },
    { :name => :e_ip,       :template => "S",  :default => 0 },
    { :name => :e_cs,       :template => "S",  :default => 0 },
    { :name => :e_lfarlc,   :template => "S",  :default => 0 },
    { :name => :e_ovno,     :template => "S",  :default => 0 },
    { :name => :e_res,      :template => "S*", :default => [0] * 4 },
    { :name => :e_oemid,    :template => "S",  :default => 0 },
    { :name => :e_oeminfo,  :template => "S",  :default => 0 },
    { :name => :e_res2,     :template => "S*", :default => [0] * 10 },
    { :name => :e_lfanew,   :template => "I",  :default => 0 },
  ],
  "ImageFileHeader" => [
    { :name => :Machine,              :template => "S", :default => 0 },
    { :name => :NumberOfSections,     :template => "S", :default => 0 },
    { :name => :TimeDateStamp,        :template => "I", :default => 0 },
    { :name => :PointerToSymbolTable, :template => "I", :default => 0 },
    { :name => :NumberOfSymbols,      :template => "I", :default => 0 },
    { :name => :SizeOfOptionalHeader, :template => "S", :default => 0 },
    { :name => :Characteristics,      :template => "S", :default => 0 },
  ],
  "ImageDataDictionary" => [
    { :name => :VirtualAddress, :template => "I", :default => 0 },
    { :name => :Size,           :template => "I", :default => 0 },
  ],
  "ImageOptionalHeader" => [
    { :name => :Magic,                       :template => "S",  :default => 0 },
    { :name => :MajorLinkerVersion,          :template => "C",  :default => 0 },
    { :name => :MinorLinkerVersion,          :template => "C",  :default => 0 },
    { :name => :SizeOfCode,                  :template => "I",  :default => 0 },
    { :name => :SizeOfInitializedData,       :template => "I",  :default => 0 },
    { :name => :SizeOfUninitializedData,     :template => "I",  :default => 0 },
    { :name => :AddressOfEntryPoint,         :template => "I",  :default => 0 },
    { :name => :BaseOfCode,                  :template => "I",  :default => 0 },
    { :name => :BaseOfData,                  :template => "I",  :default => 0 },
    { :name => :ImageBase,                   :template => "I",  :default => 0 },
    { :name => :SectionAlignment,            :template => "I",  :default => 0 },
    { :name => :FileAlignment,               :template => "I",  :default => 0 },
    { :name => :MajorOperatingSystemVersion, :template => "S",  :default => 0 },
    { :name => :MinorOperatingSystemVersion, :template => "S",  :default => 0 },
    { :name => :MajorImageVersion,           :template => "S",  :default => 0 },
    { :name => :MinorImageVersion,           :template => "S",  :default => 0 },
    { :name => :MajorSubsystemVersion,       :template => "S",  :default => 0 },
    { :name => :MinorSubsystemVersion,       :template => "S",  :default => 0 },
    { :name => :Win32VersionValue,           :template => "I",  :default => 0 },
    { :name => :SizeOfImage,                 :template => "I",  :default => 0 },
    { :name => :SizeOfHeaders,               :template => "I",  :default => 0 },
    { :name => :CheckSum,                    :template => "I",  :default => 0 },
    { :name => :Subsystem,                   :template => "S",  :default => 0 },
    { :name => :DllCharacteristics,          :template => "S",  :default => 0 },
    { :name => :SizeOfStackReserve,          :template => "I",  :default => 0 },
    { :name => :SizeOfStackCommit,           :template => "I",  :default => 0 },
    { :name => :SizeOfHeapReserve,           :template => "I",  :default => 0 },
    { :name => :SizeOfHeapCommit,            :template => "I",  :default => 0 },
    { :name => :LoaderFlags,                 :template => "I",  :default => 0 },
    { :name => :NumberOfRvaAndSizes,         :template => "I",  :default => 0 },
    { :name => :DataDirectory,               :template => "C*", :default => lambda { [ Struct::ImageDataDictionary.new ] * 16 } },
  ],
  "ImageNtHeaders32" => [
    { :name => :Signature,      :template => "a4", :default => "" },
    { :name => :FileHeader,     :template => "C*", :default => lambda { Struct::ImageFileHeader.new } },
    { :name => :OptionalHeader, :template => "C*", :default => lambda { Struct::ImageOptionalHeader.new } },
  ],
  # FIXME: Misc は 実際は VirtualSize と PhysicalAddress の union で定義されている
  "ImageSectionHeader" => [
    { :name => :Name,                 :template => "a8", :default => "" },
    { :name => :Misc,                 :template => "I",  :default => 0 },
    { :name => :VirtualAddress,       :template => "I",  :default => 0 },
    { :name => :SizeOfRawData,        :template => "I",  :default => 0 },
    { :name => :PointerToRawData,     :template => "I",  :default => 0 },
    { :name => :PointerToRelocations, :template => "I",  :default => 0 },
    { :name => :PointerToLinenumbers, :template => "I",  :default => 0 },
    { :name => :NumberOfRelocations,  :template => "S",  :default => 0 },
    { :name => :NumberOfLinenumbers,  :template => "S",  :default => 0 },
    { :name => :Characteristics,      :template => "I",  :default => 0 },
  ]
}

# 構造体の定義
structs.each do |name, fields|
  new_struct = Struct::new(name, *fields.map{|field| field[:name]})
  new_struct.fields = fields
end

# YAML 中で長さを求めるための構造体
#
# 配列で渡された各要素の length の合計を返す
# （例） !length [ "abc", "def" ] # => 6
YAML::add_domain_type("yaml.org,2002", "length") do |type, val|
  val.inject(0) {|result, item| result + item.length }
end

# YAML 中にコードを記載するための構造体
#
# コンパイル済みのバイトコードを返す
YAML::add_domain_type("yaml.org,2002", "compile") do |type, val|
  compile val.split("\n")
end

# セクションごとのアラインメントを行うための構造体
# ファイル（セクタ）のアラインメント: 512 bytes（0x200）ごとにアラインメントを行う
YAML::add_domain_type("yaml.org,2002", "section") do |type, val|
  size = val.raw.length
  padding_size = -size % 0x200
  val = [ val ] unless val.is_a? Array
  val << "\0" * padding_size
end

# バイトコードへのコンパイルを行う
def compile(lines)
  lines << "exit" unless lines.last == "exit"

  ret = ""
  lines.each_with_index do |line, i|
    line.strip!
    case line
    when /^let eax (.+)$/
      ret << [0xb8, $1.to_i].pack("CI")
    when /^exit$/
      ret << [0xc3].pack("C")
    else
      puts "can't compile(line #{i+1}): #{line}"
      exit 1
    end
  end
  ret
end

# YAML の読み込み
data = YAML.load(File.read "source.yml")
#pp data
#pp data.raw

# PE の書き出し
File.open("output.exe", "wb").write(data.raw)

source.yml

- !section
  - !ruby/struct:ImageDosHeader
    e_magic:    "MZ"
    e_cblp:     0x90
    e_cp:       3
    e_cparhdr:  4
    e_maxalloc: 0xffff
    e_sp:       0xb8
    e_lfarlc:   0x40
    e_lfanew:   0x80
  - "\xb8\x01\x4c\xcd\x21\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\
     \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\
     \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\
     \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
  - !ruby/struct:ImageNtHeaders32
    Signature: "PE"
    FileHeader: !ruby/struct:ImageFileHeader
      Machine:              0x14c
      NumberOfSections:     1
      SizeOfOptionalHeader: 0xe0
      Characteristics:      0x102
    OptionalHeader: !ruby/struct:ImageOptionalHeader
      Magic:                       0x10b
      MajorLinkerVersion:          10
      SizeOfCode:                  0x200
      AddressOfEntryPoint:         0x1000
      BaseOfCode:                  0x1000
      BaseOfData:                  0x2000
      ImageBase:                   0x400000
      SectionAlignment:            0x1000
      FileAlignment:               0x200
      MajorOperatingSystemVersion: 0x05
      MinorOperatingSystemVersion: 0x01
      MajorSubsystemVersion:       0x05
      MinorSubsystemVersion:       0x01
      SizeOfImage:                 0x2000
      SizeOfHeaders:               0x200
      Subsystem:                   2
      DllCharacteristics:          0x8500
      SizeOfStackReserve:          0x100000
      SizeOfStackCommit:           0x1000
      SizeOfHeapReserve:           0x100000
      SizeOfHeapCommit:            0x1000
      NumberOfRvaAndSizes:         0x10
  - !ruby/struct:ImageSectionHeader
    Name:             ".text"
    Misc:             !length [ &code !compile |-
      let eax 123
      exit
    ]
    VirtualAddress:   0x1000
    SizeOfRawData:    0x200
    PointerToRawData: 0x200
    Characteristics:  0x60000020
- !section
  - *code