Refactor/shellcraft (#49)

* Refactor shellcraft ( closed #46 )
* All shellcode generators are defined under `Shellcraft::Generators::${arch}::<Common|${os}>`
* The instance `shellcraft` defines `method_missing` to find the needed method according to current context.
* No longer support calls like `shellcraft.i386.linux.sh` when `context.arch` is amd64.
  - instead, use `context.local(arch: 'i386') { shellcraft.sh }` if do need that one.
  - or, use `Shellcraft::Generators::I386::Linux.sh`.
* Support `setregs` to accept registers' name in `Symbol`. ( fixed #50 )
* Fix #51

.codeclimate.yml

@@ -11,8 +11,7 @@ engines:
       - python
       - php
     exclude_paths:
-      - lib/pwnlib/shellcraft/templates/i386/mov.rb
-      - lib/pwnlib/shellcraft/templates/amd64/mov.rb
+      - lib/pwnlib/shellcraft/generators/
   fixme:
     enabled: false
   rubocop:

.rubocop.yml

@@ -48,6 +48,12 @@ Style/AutoResourceCleanup:
 Style/CollectionMethods:
   Enabled: true
 
+Style/Documentation:
+  Enabled: true
+  Exclude:
+    - 'lib/pwnlib/shellcraft/generators/**/*.rb'
+    - 'test/**/*.rb'
+
 Style/Encoding:
   Enabled: false
 

Gemfile.lock

@@ -17,7 +17,7 @@ PATH
   specs:
     pwntools (0.1.0)
       dentaku (~> 2.0.11)
-      elftools (~> 0.2.2)
+      elftools (~> 1.0.1)
       rainbow (~> 2.2)
 
 GEM
@@ -30,7 +30,7 @@ GEM
     coderay (1.1.1)
     dentaku (2.0.11)
     docile (1.1.5)
-    elftools (0.2.2)
+    elftools (1.0.1)
       bindata (~> 2)
     ffi (1.9.18)
     json (2.1.0)

lib/pwnlib/pwn.rb

@@ -32,6 +32,6 @@ module Pwn
   include ::Pwnlib::Util::Packing
 
   def shellcraft
-    ::Pwnlib::Shellcraft.instance
+    ::Pwnlib::Shellcraft::Shellcraft.instance
   end
 end

lib/pwnlib/shellcraft/generators/amd64/common/common.rb

@@ -0,0 +1,14 @@
+require 'pwnlib/shellcraft/generators/helper'
+
+module Pwnlib
+  module Shellcraft
+    module Generators
+      module Amd64
+        # For non os-related methods.
+        module Common
+          extend ::Pwnlib::Shellcraft::Generators::Helper
+        end
+      end
+    end
+  end
+end

lib/pwnlib/shellcraft/generators/amd64/common/infloop.rb

@@ -0,0 +1,17 @@
+require 'pwnlib/shellcraft/generators/amd64/common/common'
+require 'pwnlib/shellcraft/generators/x86/common/infloop'
+
+module Pwnlib
+  module Shellcraft
+    module Generators
+      module Amd64
+        module Common
+          # See {X86::Common#infloop}.
+          def infloop
+            cat Generators::X86::Common.infloop
+          end
+        end
+      end
+    end
+  end
+end

lib/pwnlib/shellcraft/generators/amd64/common/memcpy.rb

@@ -0,0 +1,31 @@
+# encoding: ASCII-8BIT
+
+require 'pwnlib/shellcraft/generators/amd64/common/common'
+require 'pwnlib/shellcraft/generators/amd64/common/setregs'
+
+module Pwnlib
+  module Shellcraft
+    module Generators
+      module Amd64
+        module Common
+          # Like +memcpy+ in glibc.
+          #
+          # Copy +n+ bytes from +src+ to +dst+.
+          #
+          # @param [String, Symbol, Integer] dst
+          #   Destination.
+          # @param [String, Symbol, Integer] src
+          #   Source to be copied.
+          # @param [Integer] n
+          #   The number of bytes to be copied.
+          def memcpy(dst, src, n)
+            cat "/* memcpy(#{pretty(dst)}, #{pretty(src)}, #{pretty(n)}) */"
+            cat 'cld'
+            cat Common.setregs(rdi: dst, rsi: src, rcx: n)
+            cat 'rep movsb'
+          end
+        end
+      end
+    end
+  end
+end

lib/pwnlib/shellcraft/generators/amd64/common/mov.rb

@@ -0,0 +1,127 @@
+# encoding: ASCII-8BIT
+
+require 'pwnlib/shellcraft/generators/amd64/common/common'
+
+module Pwnlib
+  module Shellcraft
+    module Generators
+      module Amd64
+        module Common
+          # Move +src+ into +dst+ without newlines and null bytes.
+          #
+          # @param [String, Symbol] dst
+          #   Register's name.
+          # @param [String, Symbol, Integer] src
+          #   Register's name or immediate value.
+          # @param [Boolean] stack_allowed
+          #   If equals to +false+, generated assembly code would not use stack-related operations.
+          #   But beware of without stack-related operations the generated code length is longer.
+          #
+          # @example
+          #   context.arch = 'amd64'
+          #   shellcraft.mov('rdi', 'ax')
+          #   #=> "  movzx edi, ax\n"
+          # @example
+          #   context.arch = 'amd64'
+          #   puts shellcraft.mov('rax', 10)
+          #   #   push 9 /* mov eax, '\n' */
+          #   #   pop rax
+          #   #   inc eax
+          #   #=> nil
+          # @example
+          #   context.arch = 'amd64'
+          #   puts shellcraft.mov('rax', 10, stack_allowed: false)
+          #   #   mov eax, 0x1010101
+          #   #   xor eax, 0x101010b /* 0xa == 0x1010101 ^ 0x101010b */
+          #   #=> nil
+          def mov(dst, src, stack_allowed: true)
+            raise ArgumentError, "#{dst} is not a register" unless register?(dst)
+            dst = get_register(dst)
+            if register?(src)
+              src = get_register(src)
+              if dst.size < src.size && !dst.bigger.include?(src.name)
+                raise ArgumentError, "cannot mov #{dst}, #{src}: dst is smaller than src"
+              end
+              # Downgrade our register choice if possible.
+              # Opcodes for operating on 32-bit registers are always (?) shorter.
+              dst = get_register(dst.native32) if dst.size == 64 && src.size <= 32
+            else
+              context.local(arch: 'amd64') { src = evaluate(src) }
+              raise ArgumentError, format('cannot mov %s, %d: dst is smaller than src', dst, src) unless dst.fits(src)
+              orig_dst = dst
+              dst = get_register(dst.native32) if dst.size == 64 && bits_required(src) <= 32
+
+              # Calculate the packed version.
+              srcp = pack(src & ((1 << dst.size) - 1), bits: dst.size)
+
+              # Calculate the unsigned and signed versions.
+              srcu = unpack(srcp, bits: dst.size, signed: false)
+              # N.B.: We may have downsized the register for e.g. mov('rax', 0xffffffff)
+              #       In this case, srcp is now a 4-byte packed value, which will expand to "-1", which isn't correct.
+              srcs = orig_dst.size == dst.size ? unpack(srcp, bits: dst.size, signed: true) : src
+            end
+            if register?(src)
+              if src == dst || dst.bigger.include?(src.name)
+                cat "/* moving #{src} into #{dst}, but this is a no-op */"
+              elsif dst.size > src.size
+                cat "movzx #{dst}, #{src}"
+              else
+                cat "mov #{dst}, #{src}"
+              end
+            elsif src.is_a?(Numeric) # Constant or immi
+              xor = ->(reg) { "xor #{reg.xor}, #{reg.xor}" }
+              if src.zero?
+                # Special case for zeroes.
+                # XORing the 32-bit register clears the high 32 bits as well.
+                cat "xor #{dst}, #{dst} /* #{src} */"
+              elsif stack_allowed && [32, 64].include?(dst.size) && src == 10
+                cat "push 9 /* mov #{dst}, '\\n' */"
+                cat "pop #{dst.native64}"
+                cat "inc #{dst}"
+              elsif stack_allowed && [32, 64].include?(dst.size) && (-2**7 <= srcs && srcs < 2**7) && okay(srcp[0])
+                # It's smaller to PUSH and POP small sign-extended values than to directly move them into various
+                # registers.
+                #
+                # 6aff58           push -1; pop rax
+                # 48c7c0ffffffff   mov rax, -1
+                cat "push #{pretty(src)}"
+                cat "pop #{dst.native64}"
+              elsif okay(srcp)
+                # Easy case. This implies that the register size and value are the same.
+                cat "mov #{dst}, #{pretty(src)}"
+              elsif srcu < 2**8 && okay(srcp[0]) && dst.sizes.include?(8) # Move 8-bit value into register.
+                cat xor[dst]
+                cat "mov #{dst.sizes[8]}, #{pretty(src)}"
+              elsif srcu == srcu & 0xff00 && okay(srcp[1]) && dst.ff00
+                # Target value is a 16-bit value with no data in the low 8 bits, we can use the 'AH' style register.
+                cat xor[dst]
+                cat "mov #{dst.ff00}, #{pretty(src)} >> 8"
+              elsif srcu < 2**16 && okay(srcp[0, 2]) # Target value is a 16-bit value, use a 16-bit mov.
+                cat xor[dst]
+                cat "mov #{dst.sizes[16]}, #{pretty(src)}"
+              else # All else has failed.  Use some XOR magic to move things around.
+                a, b = xor_pair(srcp, avoid: "\x00\n")
+                a = hex(unpack(a, bits: dst.size))
+                b = hex(unpack(b, bits: dst.size))
+                if dst.size != 64
+                  # There's no XOR REG, IMM64 but we can take the easy route for smaller registers.
+                  cat "mov #{dst}, #{a}"
+                  cat "xor #{dst}, #{b} /* #{hex(src)} == #{a} ^ #{b} */"
+                elsif stack_allowed
+                  # However, we can PUSH IMM64 and then perform the XOR that way at the top of the stack.
+                  cat "mov #{dst}, #{a}"
+                  cat "push #{dst}"
+                  cat "mov #{dst}, #{b}"
+                  cat "xor [rsp], #{dst} /* #{hex(src)} == #{a} ^ #{b} */"
+                  cat "pop #{dst}"
+                else
+                  raise ArgumentError, "Cannot put #{pretty(src)} into '#{dst}' without using stack."
+                end
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/pwnlib/shellcraft/generators/amd64/common/nop.rb

@@ -0,0 +1,16 @@
+require 'pwnlib/shellcraft/generators/amd64/common/common'
+
+module Pwnlib
+  module Shellcraft
+    module Generators
+      module Amd64
+        module Common
+          # A no-op instruction.
+          def nop
+            cat 'nop'
+          end
+        end
+      end
+    end
+  end
+end

lib/pwnlib/shellcraft/generators/amd64/common/popad.rb

@@ -0,0 +1,27 @@
+# encoding: ASCII-8BIT
+
+require 'pwnlib/shellcraft/generators/amd64/common/common'
+
+module Pwnlib
+  module Shellcraft
+    module Generators
+      module Amd64
+        module Common
+          # Pop all of the registers onto the stack which i386 +popad+ does.
+          def popad
+            cat <<-EOS
+pop rdi
+pop rsi
+pop rbp
+pop rbx /* add rsp, 8 */
+pop rbx
+pop rdx
+pop rcx
+pop rax
+            EOS
+          end
+        end
+      end
+    end
+  end
+end

lib/pwnlib/shellcraft/generators/amd64/common/pushstr.rb

@@ -0,0 +1,64 @@
+# encoding: ASCII-8BIT
+
+require 'pwnlib/shellcraft/generators/amd64/common/common'
+
+module Pwnlib
+  module Shellcraft
+    module Generators
+      module Amd64
+        module Common
+          # Push a string to stack.
+          #
+          # @param [String] str
+          #   String to be pushed.
+          # @param [Boolean] append_null
+          #   If need to append a null byte in the end of +str+.
+          #
+          # @example
+          #   context.arch = 'amd64'
+          #   puts shellcraft.pushstr('pusheen')
+          #   #   /* push "pusheen\x00" */
+          #   #   mov rax, 0x101010101010101
+          #   #   push rax
+          #   #   mov rax, 0x101010101010101 ^ 0x6e656568737570
+          #   #   xor [rsp], rax
+          #   #=> nil
+          def pushstr(str, append_null: true)
+            # This will not affect callee's +str+.
+            str += "\x00" if append_null && !str.end_with?("\x00")
+            return if str.empty?
+            padding = str[-1].ord >= 128 ? "\xff" : "\x00"
+            cat "/* push #{str.inspect} */"
+            group(8, str, underfull_action: :fill, fill_value: padding).reverse_each do |word|
+              sign = u64(word, endian: 'little', signed: true)
+              sign32 = u32(word[0, 4], bits: 32, endian: 'little', signed: true)
+              if [0, 0xa].include?(sign) # simple forbidden byte case
+                cat "push #{pretty(sign + 1)}"
+                cat 'dec byte ptr [rsp]'
+              elsif sign >= -0x80 && sign <= 0x7f && okay(word[0]) # simple byte case
+                cat "push #{pretty(sign)}"
+              elsif sign >= -0x80000000 && sign <= 0x7fffffff && okay(word[0, 4])
+                # simple 32bit without forbidden byte
+                cat "push #{pretty(sign)}"
+              elsif okay(word)
+                cat "mov rax, #{pretty(sign)}"
+                cat 'push rax'
+              elsif sign32 > 0 && word[4, 4] == "\x00" * 4
+                # The high 4 byte of word are all zeros, so we can use +xor dword ptr [rsp]+.
+                a = u32(xor_pair(word[0, 4]).first, endian: 'little', signed: true)
+                cat "push #{pretty(a)} ^ #{pretty(sign)}"
+                cat "xor dword ptr [rsp], #{pretty(a)}"
+              else
+                a = u64(xor_pair(word).first, endian: 'little', signed: false)
+                cat "mov rax, #{pretty(a)}"
+                cat 'push rax'
+                cat "mov rax, #{pretty(a ^ sign)} /* #{pretty(a)} ^ #{pretty(sign)} */"
+                cat 'xor [rsp], rax'
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/pwnlib/shellcraft/generators/amd64/common/pushstr_array.rb

@@ -0,0 +1,19 @@
+require 'pwnlib/shellcraft/generators/amd64/common/common'
+require 'pwnlib/shellcraft/generators/x86/common/pushstr_array'
+
+module Pwnlib
+  module Shellcraft
+    module Generators
+      module Amd64
+        module Common
+          # See {Pwnlib::Shellcraft::Generators::X86::Common#pushstr_array}.
+          def pushstr_array(*args)
+            context.local(arch: 'amd64') do
+              cat X86::Common.pushstr_array(*args)
+            end
+          end
+        end
+      end
+    end
+  end
+end