Multiple enhancements to shellcraft

- Refactor i386/amd64 mov code to use new register info
	- Mostly this is to clean up the code and remove things like bigger[0][-1]
- A small enhancement to the ARM shellcode
- Resolve hosts for i386.linux.connect
- Add arm.trap
- Add gethostbyname to i386.connect

pwnlib/constants/__init__.py

@@ -87,9 +87,13 @@ class module(ModuleType):
             ...    print constants.eval('SYS_execve + PROT_WRITE')
             61
         """
+        if not isinstance(string, str):
+            return string
+
         key = context.os, context.arch
         if key not in self._env_store:
             self._env_store[key] = {key: getattr(self, key) for key in dir(self) if not key.endswith('__')}
+
         return safeeval.values(string, self._env_store[key])
 
 # To prevent garbage collection

pwnlib/shellcraft/__init__.py

@@ -1,8 +1,13 @@
+import os
+import re
+import sys
 from types import ModuleType
-import sys, os, re
+
 from . import internal
+from .. import constants
 from ..context import context
 
+
 class module(ModuleType):
     def __init__(self, name, directory):
         super(module, self).__init__(name)
@@ -112,6 +117,10 @@ class module(ModuleType):
 
     templates = sorted(templates)
 
+    def eval(self, item):
+        return constants.eval(item)
+
+    import registers
 
 # To prevent garbage collection
 tether = sys.modules[__name__]

pwnlib/shellcraft/internal.py

@@ -75,6 +75,28 @@ def lookup_template(filename):
 
     return loaded[filename]
 
+def get_context_from_dirpath(directory):
+    """
+    >>> get_context_from_dirpath('common')
+    {}
+    >>> get_context_from_dirpath('i386')
+    {'arch': 'i386'}
+    >>> get_context_from_dirpath('amd64/linux') == {'arch': 'amd64', 'os': 'linux'}
+    True
+    """
+    A,O = os.path.split(directory)
+
+    if O == 'common':
+        O = None
+
+    if not A:
+        A,O = O,None
+
+    rv = {}
+    if O: rv['os']=O
+    if A: rv['arch']=A
+    return rv
+
 def make_function(funcname, filename, directory):
     import inspect
     path       = os.path.join(directory, filename)
@@ -100,19 +122,23 @@ def make_function(funcname, filename, directory):
         args.append('**' + keywords)
         args_used.append('**' + keywords)
 
+    docstring = inspect.cleandoc(template.module.__doc__ or '')
     args      = ', '.join(args)
     args_used = ', '.join(args_used)
+    local_ctx = get_context_from_dirpath(directory)
 
     # This is a slight hack to get the right signature for the function
     # It would be possible to simply create an (*args, **kwargs) wrapper,
     # but what would not have the right signature.
     # While we are at it, we insert the docstring too
-    exec '''
+    T = '''
 def wrap(template, render_global):
-    def %s(%s):
-        %r
+    import pwnlib
+    def %(funcname)s(%(args)s):
+        %(docstring)r
         with render_global.go_inside() as was_inside:
-            lines = template.render(%s).split('\\n')
+            with pwnlib.context.context.local(**%(local_ctx)s):
+                lines = template.render(%(args_used)s).split('\\n')
         for i in xrange(len(lines)):
             line = lines[i]
             def islabelchar(c):
@@ -132,8 +158,10 @@ def wrap(template, render_global):
             return s
         else:
             return s + '\\n'
-    return %s
-''' % (funcname, args, inspect.cleandoc(template.module.__doc__ or ''), args_used, funcname)
+    return %(funcname)s
+''' % locals()
+
+    exec T
 
     # Setting _relpath is a slight hack only used to get better documentation
     res = wrap(template, render_global)

pwnlib/shellcraft/registers.py

@@ -0,0 +1,166 @@
+import re
+
+from ..context import context
+from ..util.misc import register_sizes
+
+mips =  map('r{}'.format, range(32))
+mips += map('v{}'.format, range(2))
+mips += map('a{}'.format, range(4))
+mips += map('t{}'.format, range(8))
+mips += map('s{}'.format, range(9))
+mips += map('t{}'.format, range(8,10))
+mips += map('k{}'.format, range(2))
+mips += ['zero', 'at', 'gp', 'sp', 'ra']
+
+arm = map('r{}'.format, range(13))
+arm += ["sp", "lr", "pc", "cpsr"]
+
+i386_baseregs = [ "ax", "cx", "dx", "bx", "sp", "bp", "si", "di", "ip"]
+
+i386 = map('e{}'.format, i386_baseregs)
+i386 += i386_baseregs
+i386 += [ "eflags", "cs", "ss", "ds", "es", "fs", "gs", ]
+
+amd64 =  map('r{}'.format, i386_baseregs)
+amd64 += map('r{}'.format, range(10,16))
+amd64 += map('r{}d'.format, range(10,16))
+amd64 += i386
+
+powerpc =  map('r{}'.format, range(32))
+powerpc += ["pc", "msr", "cr", "lr", "ctr", "xer", "orig_r3", "trap" ]
+powerpc =  map('%{}'.format, powerpc)
+
+sparc =  map('g{}'.format, range(8))
+sparc += map('o{}'.format, range(5))
+sparc += map('l{}'.format, range(8))
+sparc += map('i{}'.format, range(5))
+sparc += ["pc", "sp", "fp", "psr" ]
+sparc =  map('%{}'.format, sparc)
+
+
+
+# x86/amd64 registers in decreasing size
+i386_ordered = [
+    ['rax', 'eax', 'ax', 'al'],
+    ['rbx', 'ebx', 'bx', 'bl'],
+    ['rcx', 'ecx', 'cx', 'cl'],
+    ['rdx', 'edx', 'dx', 'dl'],
+    ['rdi', 'edi', 'di'],
+    ['rsi', 'esi', 'si'],
+    ['rbp', 'ebp', 'bp'],
+    ['rsp', 'esp', 'sp'],
+    ['r8', 'r8d', 'r8w', 'r8b'],
+    ['r9', 'r9d', 'r9w', 'r9b'],
+    ['r10', 'r10d', 'r10w', 'r10b'],
+    ['r11', 'r11d', 'r11w', 'r11b'],
+    ['r12', 'r12d', 'r12w', 'r12b'],
+    ['r13', 'r13d', 'r13w', 'r13b'],
+    ['r14', 'r14d', 'r14w', 'r14b'],
+    ['r15', 'r15d', 'r15w', 'r15b']
+]
+
+all_regs, sizes, bigger, smaller = register_sizes(i386_ordered, [64, 32, 16, 8, 8])
+native64 = {k:v[0] for k,v in bigger.items()}
+native32 = {k:v[1] for k,v in bigger.items() if not k.startswith('r')}
+
+class Register(object):
+    #: Register name
+    name = None
+
+    #: List of larger registers, in order from largest to smallest
+    bigger = None
+
+    #: List of smaller regsters, in order from smallest to largest
+    smaller = None
+
+    #: Size of the register, in bits
+    size = None
+
+    #: Does this register have a 'high' register for mask 0xff00
+    ff00 = None
+
+    #: Flags for 64-bit mode.64-bit
+    #: The first bit is set, if the register can be used with a REX-mode
+    #: The second bit is set, if the register can be used without a REX-prefix
+    rex_mode = 0
+
+    #: Is this a 64-bit only register?
+    is64bit = False
+
+    #: Name of the native 64-bit register
+    native64 = None
+
+    #: Name of the native 32-bit register
+    native32 = None
+
+    #: Name of the register which should be used to clear
+    #: this register, e.g. xor REG, REG.
+    #: Useful for AMD64 for xor eax, eax is shorter than
+    #: xor rax, rax and has the same effect.
+    xor = None
+
+    def __init__(self, name, size):
+        self.name = name
+        self.size = size
+
+        for row in i386_ordered:
+            if name in row:
+                self.bigger  = row[0:row.index(name)]
+                self.smaller = row[row.index(name)+1:]
+                self.sizes   = {64>>i:r for i,r in enumerate(row)}
+                self.native64 = row[0]
+                self.native32 = row[1]
+                self.xor = self.sizes[min(self.size, 32)]
+
+        if self.size >= 32 and name.endswith('x'):
+            self.ff00 = name[1] + 'h'
+
+        if name[-1] != 'h':
+            self.rex_mode |= 1
+
+        if name[0] != 'r':
+            self.rex_mode |= 2
+
+        if name.startswith('r') or name[1:3].isdigit():
+            self.is64bit = True
+
+    def fits(self, value):
+        return self.size >= bits_required(value)
+
+    def __str__(self):
+        return self.name
+
+    def __repr__(self):
+        return "Register(%r)" % self.namepyth
+
+intel = {}
+
+for row in i386_ordered:
+    for i, reg in enumerate(row):
+        intel[reg] = Register(reg, 64 >> i)
+
+def get_register(name):
+    return intel.get(name, None)
+
+def is_register(obj):
+    if isinstance(obj, Register):
+        return True
+    return get_register(obj)
+
+
+def bits_required(value):
+    bits  = 0
+
+    if value < 0:
+        return context.bits
+
+    while value:
+        value >>= 8
+        bits += 8
+    return bits
+
+def register_size(reg):
+    return sizes[reg]
+
+def fits_in_register(reg, value):
+    return register_size(reg) >= bits_required(value)

pwnlib/shellcraft/templates/amd64/mov.asm

@@ -3,6 +3,8 @@
   from pwnlib import constants
   from pwnlib.context import context as ctx # Ugly hack, mako will not let it be called context
   from pwnlib.log import getLogger
+  from pwnlib.shellcraft import i386
+  from pwnlib.shellcraft.registers import get_register, is_register, bits_required
   log = getLogger('pwnlib.shellcraft.amd64.mov')
 %>
 <%page args="dest, src, stack_allowed = True"/>
@@ -30,10 +32,10 @@ Example:
         xor ax, ax
     >>> print shellcraft.amd64.mov('rax', 0).rstrip()
         xor eax, eax
+    >>> print shellcraft.amd64.mov('rdi', 'ax').rstrip()
+        movzx edi, ax
     >>> print shellcraft.amd64.mov('al', 'ax').rstrip()
         /* moving ax into al, but this is a no-op */
-    >>> print shellcraft.amd64.mov('bl', 'ax').rstrip()
-        mov bl, al
     >>> print shellcraft.amd64.mov('ax', 'bl').rstrip()
         movzx ax, bl
     >>> print shellcraft.amd64.mov('eax', 1).rstrip()
@@ -51,12 +53,6 @@ Example:
     >>> print shellcraft.amd64.mov('rax', 0xdead00ff).rstrip()
         mov eax, 0x1010101
         xor eax, 0xdfac01fe
-    >>> print shellcraft.amd64.mov('ah', 'al').rstrip()
-        mov ah, al
-    >>> print shellcraft.amd64.mov('ah', 'r8').rstrip()
-    Traceback (most recent call last):
-    ...
-    PwnlibException: The amd64 instruction set does not support moving from r8 to ah
     >>> print shellcraft.amd64.mov('rax', 0x11dead00ff).rstrip()
         mov rax, 0x101010101010101
         push rax
@@ -81,99 +77,133 @@ Args:
   stack_allowed (bool): Can the stack be used?
 </%docstring>
 <%
-regs = [['rax', 'eax', 'ax', 'ah', 'al'],
-        ['rbx', 'ebx', 'bx', 'bh', 'bl'],
-        ['rcx', 'ecx', 'cx', 'ch', 'cl'],
-        ['rdx', 'edx', 'dx', 'dh', 'dl'],
-        ['rdi', 'edi', 'di', 'dil'],
-        ['rsi', 'esi', 'si', 'sil'],
-        ['rbp', 'ebp', 'bp', 'bpl'],
-        ['rsp', 'esp', 'sp', 'spl'],
-        ['r8', 'r8d', 'r8w', 'r8b'],
-        ['r9', 'r9d', 'r9w', 'r9b'],
-        ['r10', 'r10d', 'r10w', 'r10b'],
-        ['r11', 'r11d', 'r11w', 'r11b'],
-        ['r12', 'r12d', 'r12w', 'r12b'],
-        ['r13', 'r13d', 'r13w', 'r13b'],
-        ['r14', 'r14d', 'r14w', 'r14b'],
-        ['r15', 'r15d', 'r15w', 'r15b']
-        ]
 def okay(s):
     return '\0' not in s and '\n' not in s
 
 def pretty(n):
     if n < 0:
-      return str(n)
+        return str(n)
     else:
-      return hex(n)
+        return hex(n)
 
-all_regs, sizes, bigger, smaller = misc.register_sizes(regs, [64, 32, 16, 8, 8])
+src_name = src
+if not isinstance(src, (str, tuple)):
+    src_name = pretty(src)
 
+if not get_register(dest):
+    log.error('%r is not a register' % dest)
 
-if isinstance(src, (str, unicode)):
-    src = src.strip()
-    if src.lower() in all_regs:
-        src = src.lower()
-    else:
-        with ctx.local(arch = 'amd64'):
-            try:
-                src = constants.eval(src)
-            except:
-                log.error("Could not figure out the value of %r" % src)
+dest = get_register(dest)
 
-dest_orig = dest
-%>
-% if dest not in all_regs:
-    <% log.error('%s is not a register' % str(dest_orig)) %>\
+if get_register(src):
+    src = get_register(src)
+
+    if dest.size < src.size and src.name not in dest.bigger:
+        log.error("cannot mov %s, %s: dddest is smaller than src" % (dest, src))
+
+    # Can't move between RAX and DIL for example.
+    if dest.rex_mode & src.rex_mode == 0:
+        log.error('The amd64 instruction set does not support moving from %s to %s' % (src, dest))
+
+    # Downgrade our register choice if possible.
+    # Opcodes for operating on 32-bit registers are always (?) shorter.
+    if dest.size == 64 and src.size <= 32:
+        dest = get_register(dest.native32)
+
+    src_size = src.size
+else:
+    with ctx.local(arch = 'amd64'):
+        src = constants.eval(src)
+
+    if not dest.fits(src):
+        log.error("cannot mov %s, %r: dest is smaller than src" % (dest, src))
+
+    src_size = bits_required(src)
+
+    if dest.size == 64 and src_size <= 32:
+        dest = get_register(dest.native32)
+
+    # Calculate the packed version
+    srcp = packing.pack(src & ((1<<64)-1), dest.size)
+
+    # Calculate the unsigned and signed versions
+    srcu = packing.unpack(srcp, dest.size, sign=False)
+    srcs = packing.unpack(srcp, dest.size, sign=True)
+
+%>\
+% if is_register(src):
+    % if src == dest:
+    /* moving ${src} into ${dest}, but this is a no-op */
+    % elif src.name in dest.bigger:
+    /* moving ${src} into ${dest}, but this is a no-op */
+    % elif dest.size > src.size:
+    movzx ${dest}, ${src}
+    % else:
+    mov ${dest}, ${src}
+    % endif
 % elif isinstance(src, (int, long)):
-    <%
-      if not (-2 ** (sizes[dest]-1) <= src < 2**sizes[dest]):
-          log.error('Number %s does not fit into %s' % (pretty(src), dest_orig))
-
-      # Calculate the unsigned and signed versions
-      srcu = src & (2 ** sizes[dest] - 1)
-      srcs = srcu - 2 * (srcu & (2 ** (sizes[dest] - 1)))
-
-      # micro-optimization: if you ever touch e.g. eax, then all the upper bits
-      # of rax will be set to 0. We exploit this fact here
-      if 0 <= src < 2 ** 32 and sizes[dest] == 64:
-          dest = smaller[dest][0]
-
-      # Calculate the packed version
-      srcp = packing.pack(srcu, sizes[dest], 'little', False)
-    %>\
+## Special case for zeroes
+## XORing the 32-bit register clears the high 32 bits as well
     % if src == 0:
         xor ${dest}, ${dest}
-    % elif src == 10 and stack_allowed and sizes[dest] == 32: # sizes[dest] == 64 not possible here
-        push 9
-        pop ${bigger[dest][0]}
+## Special case for *just* a newline
+    % elif stack_allowed and dest.size in (32,64) and src == 10:
+        push 9 /* mov ${dest}, '\n' */
+        pop ${dest.native64}
         inc ${dest}
-    % elif stack_allowed and sizes[dest] in [32, 64] and -2**7 <= srcs < 2**7 and 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
+    % elif stack_allowed and dest.size in (32,64) and (-2**7 <= srcs < 2**7) and okay(srcp[:1]):
         push ${pretty(srcs)}
-        pop ${dest if sizes[dest] == 64 else bigger[dest][0]}
+        pop ${dest.native64}
+## Easy case, everybody is trivially happy
+## This implies that the register size and value are the same.
     % elif okay(srcp):
         mov ${dest}, ${pretty(src)}
-    % elif stack_allowed and sizes[dest] in [32, 64] and -2**31 <= srcs < 2**31 and okay(srcp[:4]):
+## We can push 32-bit values onto the stack and they are sign-extended.
+    % elif stack_allowed and dest.size in (32,64) and (-2**31 <= srcs < 2**31) and okay(srcp[:4]):
         push ${pretty(srcs)}
-        pop ${dest if sizes[dest] == 64 else bigger[dest][0]}
-    % elif 0 <= srcu < 2**8 and okay(srcp[0]) and sizes[smaller[dest][-1]] == 8:
+        pop ${dest.native64}
+## We can also leverage the sign-extension to our advantage.
+## For example, 0xdeadbeef is sign-extended to 0xffffffffdeadbeef.
+## Want EAX=0xdeadbeef, we don't care that RAX=0xfff...deadbeef.
+    % elif stack_allowed and dest.size == 32 and srcu < 2**32 and okay(srcp[:4]):
+        push ${pretty(srcs)}
+        pop ${dest.native64}
+## Target value is an 8-bit value, use a 8-bit mov
+    % elif srcu < 2**8 and okay(srcp[:1]):
+        xor ${dest.xor}, ${dest.xor}
+        mov ${dest.sizes[8]}, ${pretty(srcu)}
+## Target value is a 16-bit value with no data in the low 8 bits
+## means we can use the 'AH' style register.
+    % elif srcu == srcu & 0xff00 and okay(srcp[1]) and dest.ff00:
         xor ${dest}, ${dest}
-        mov ${smaller[dest][-1]}, ${pretty(srcu)}
-    % elif srcu == srcu & 0xff00 and okay(srcp[1]) and sizes[smaller[dest][-2]] == 8:
-        xor ${dest}, ${dest}
-        mov ${smaller[dest][-2]}, ${pretty(srcu >> 8)}
-    % elif 0 <= srcu < 2**16 and okay(srcp[:2]) and sizes[smaller[dest][0]] == 16:
-        xor ${dest}, ${dest}
-        mov ${smaller[dest][0]}, ${pretty(src)}
+        mov ${dest.ff00}, ${pretty(srcu >> 8)}
+## Target value is a 16-bit value, use a 16-bit mov
+    % elif srcu < 2**16 and okay(srcp[:2]):
+        xor ${dest.xor}, ${dest.xor}
+        mov ${dest.sizes[16]}, ${pretty(srcu)}
+## Target value is a 32-bit value, use a 32-bit mov.
+## Note that this is zero-extended rather than sign-extended (the 32-bit push above).
+    % elif srcu < 2**32 and okay(srcp[:4]):
+        mov ${dest.sizes[32]}, ${pretty(srcu)}
+## All else has failed.  Use some XOR magic to move things around.
     % else:
         <%
         a,b = fiddling.xor_pair(srcp, avoid = '\x00\n')
-        a = pretty(packing.unpack(a, sizes[dest], 'little', False))
-        b = pretty(packing.unpack(b, sizes[dest], 'little', False))
+        a = pretty(packing.unpack(a, dest.size, 'little', False))
+        b = pretty(packing.unpack(b, dest.size, 'little', False))
         %>\
-        % if sizes[dest] != 64:
+## There's no XOR REG, IMM64 but we can take the easy route
+## for smaller registers.
+        % if dest.size != 64:
           mov ${dest}, ${a}
           xor ${dest}, ${b}
+## However, we can PUSH IMM64 and then perform the XOR that
+## way at the top of the stack.
         % elif stack_allowed:
           mov ${dest}, ${a}
           push ${dest}
@@ -184,42 +214,6 @@ dest_orig = dest
           <% log.error("Cannot put %s into '%s' without using stack." % (pretty(src), dest_orig)) %>\
         % endif
     % endif
-% elif src in all_regs:
-    <%
-      # micro-optimization: if you ever touch e.g. eax, then all the upper bits
-      # of rax will be set to 0. We exploit this fact here
-      if sizes[dest] == 64 and sizes[src] != 64:
-          dest = smaller[dest][0]
-
-      def rex_mode(reg):
-          # This returns a number with two bits
-          # The first bit is set, if the register can be used with a REX-mode
-          # The seond bit is set, if the register can be used without a REX-prefix
-          res = 0
-          if reg[-1] != 'h':
-              res |= 1
-
-          if reg[0] != 'r' and reg not in ['dil', 'sil', 'bpl', 'spl']:
-              res |= 2
-
-          return res
-      if rex_mode(src) & rex_mode(dest) == 0:
-          log.error('The amd64 instruction set does not support moving from %s to %s' % (src, dest))
-    %>\
-    % if src == dest or src in bigger[dest] or src in smaller[dest]:
-        /* moving ${src} into ${dest_orig}, but this is a no-op */
-    % elif sizes[dest] == sizes[src]:
-        mov ${dest}, ${src}
-    % elif sizes[dest] > sizes[src]:
-        movzx ${dest}, ${src}
-    % else:
-        % for r in reversed(smaller[src]):
-            % if sizes[r] == sizes[dest]:
-                mov ${dest}, ${r}
-                <% return %>\
-            % endif
-        % endfor
-    % endif
 % else:
     <% log.error('%s is neither a register nor an immediate' % src) %>\
-% endif
+% endif

pwnlib/shellcraft/templates/arm/mov.asm

@@ -6,10 +6,10 @@
     Returns THUMB code for moving the specified source value
     into the specified destination register.
 </%docstring>
-/* Set ${dst} = ${src} = 0x${'%x' % src} */
 %if not isinstance(src, (int, long)):
     mov ${dst}, ${src}
 %else:
+/* Set ${dst} = ${src} = 0x${'%x' % src} */
   %if src == 0:
     eor ${dst}, ${dst}
   %elif src & 0xffff0000 == 0:

pwnlib/shellcraft/templates/arm/trap.asm

@@ -0,0 +1,2 @@
+<%docstring>A trap instruction.</%docstring>
+    BKPT

pwnlib/shellcraft/templates/i386/linux/connect.asm

@@ -1,6 +1,6 @@
 <% from pwnlib.shellcraft import common %>
-<% from pwnlib.shellcraft import i386 %> 
-<% from socket import htons, inet_aton %>
+<% from pwnlib.shellcraft import i386 %>
+<% from socket import htons, inet_aton, gethostbyname %>
 <% from pwnlib.util import packing %>
 
 <%page args="host, port"/>
@@ -24,7 +24,10 @@ int 0x80
 /* save opened socket */
 mov ebp, eax
 
-${i386.pushstr(inet_aton(host), False)}
+<% ip_addr = gethostbyname(str(host)) %>
+/* ${repr(host)} == ${ip_addr} */
+${i386.pushstr(inet_aton(ip_addr), False)}
+
 pushw ${htons(port)}
 pushw AF_INET
 mov ecx, esp

pwnlib/shellcraft/templates/i386/linux/syscall.asm

@@ -56,7 +56,7 @@ Example:
         ...               -1, 0).rstrip()
             /* call mmap2(0, 4096, 'PROT_READ | PROT_WRITE | PROT_EXEC', 'MAP_PRIVATE | MAP_ANONYMOUS', -1, 0) */
             xor ebx, ebx
-            xor ecx, ecx
+            xor ecx, ecx /* mov ecx, 0x1000 */
             mov ch, 0x10
             push 0x7
             pop edx

pwnlib/shellcraft/templates/i386/mov.asm

@@ -3,6 +3,7 @@
   from pwnlib import constants
   from pwnlib.context import context as ctx # Ugly hack, mako will not let it be called context
   from pwnlib.log import getLogger
+  from pwnlib.shellcraft.registers import get_register, is_register, bits_required
   log = getLogger('pwnlib.shellcraft.i386.mov')
 %>
 <%page args="dest, src, stack_allowed = True"/>
@@ -31,37 +32,47 @@ Example:
     >>> print shellcraft.i386.mov('ax', 17).rstrip()
         xor ax, ax
         mov al, 0x11
+    >>> print shellcraft.i386.mov('edi', ord('\n')).rstrip()
+        push 9 /* mov edi, '\n' */
+        pop edi
+        inc edi
     >>> print shellcraft.i386.mov('al', 'ax').rstrip()
         /* moving ax into al, but this is a no-op */
-    >>> print shellcraft.i386.mov('bl', 'ax').rstrip()
-        mov bl, al
+    >>> print shellcraft.i386.mov('al','ax').rstrip()
+        /* moving ax into al, but this is a no-op */
+    >>> print shellcraft.i386.mov('esp', 'esp').rstrip()
+        /* moving esp into esp, but this is a no-op */
     >>> print shellcraft.i386.mov('ax', 'bl').rstrip()
         movzx ax, bl
     >>> print shellcraft.i386.mov('eax', 1).rstrip()
         push 0x1
         pop eax
+    >>> print shellcraft.i386.mov('eax', 1, stack_allowed=False).rstrip()
+        xor eax, eax
+        mov al, 0x1
     >>> print shellcraft.i386.mov('eax', 0xdead00ff).rstrip()
-        mov eax, 0x1010101
+        mov eax, 0x1010101 /* mov eax, 0xdead00ff */
         xor eax, 0xdfac01fe
     >>> print shellcraft.i386.mov('eax', 0xc0).rstrip()
         xor eax, eax
         mov al, 0xc0
     >>> print shellcraft.i386.mov('eax', 0xc000).rstrip()
-        xor eax, eax
+        xor eax, eax /* mov eax, 0xc000 */
         mov ah, 0xc0
+    >>> print shellcraft.i386.mov('edi', 0xc000).rstrip()
+        mov edi, 0x1010101 /* mov edi, 0xc000 */
+        xor edi, 0x101c101
     >>> print shellcraft.i386.mov('eax', 0xc0c0).rstrip()
         xor eax, eax
         mov ax, 0xc0c0
-    >>> with context.local(os = 'linux'):
-    ...     print shellcraft.i386.mov('eax', 'SYS_execve').rstrip()
+    >>> print shellcraft.i386.mov('eax', 'SYS_execve').rstrip()
         push 0xb
         pop eax
     >>> with context.local(os = 'freebsd'):
     ...     print shellcraft.i386.mov('eax', 'SYS_execve').rstrip()
         push 0x3b
         pop eax
-    >>> with context.local(os = 'linux'):
-    ...     print shellcraft.i386.mov('eax', 'PROT_READ | PROT_WRITE | PROT_EXEC').rstrip()
+    >>> print shellcraft.i386.mov('eax', 'PROT_READ | PROT_WRITE | PROT_EXEC').rstrip()
         push 0x7
         pop eax
 
@@ -71,16 +82,6 @@ Args:
   stack_allowed (bool): Can the stack be used?
 </%docstring>
 <%
-regs = [['eax', 'ax', 'ah', 'al'],
-        ['ebx', 'bx', 'bh', 'bl'],
-        ['ecx', 'cx', 'ch', 'cl'],
-        ['edx', 'dx', 'dh', 'dl'],
-        ['edi', 'di'],
-        ['esi', 'si'],
-        ['ebp', 'bp'],
-        ['esp', 'sp'],
-        ]
-
 def okay(s):
     return '\0' not in s and '\n' not in s
 
@@ -90,84 +91,97 @@ def pretty(n):
     else:
         return hex(n)
 
-all_regs, sizes, bigger, smaller = misc.register_sizes(regs, [32, 16, 8, 8])
+src_name = src
+if not isinstance(src, (str, tuple)):
+    src_name = pretty(src)
 
-if isinstance(src, (str, unicode)):
-    src = src.strip()
-    if src.lower() in all_regs:
-        src = src.lower()
-    else:
-        with ctx.local(arch = 'i386'):
-            try:
-                src = constants.eval(src)
-            except (AttributeError, ValueError):
-                log.error("Could not figure out the value of %r" % src)
-                return
-%>
-% if dest not in all_regs:
-   <% log.error('%s is not a register' % str(dest)) %>\
+if not get_register(dest):
+    log.error('%r is not a register' % dest)
+
+dest = get_register(dest)
+
+if dest.size > 32 or dest.is64bit:
+    log.error("cannot use %s on i386" % dest)
+
+if get_register(src):
+    src = get_register(src)
+
+    if src.is64bit:
+        log.error("cannot use %s on i386" % src)
+
+    if dest.size < src.size and src.name not in dest.bigger:
+        log.error("cannot mov %s, %s: dddest is smaller than src" % (dest, src))
+else:
+    with ctx.local(arch = 'i386'):
+        src = constants.eval(src)
+
+    if not dest.fits(src):
+        log.error("cannot mov %s, %r: dest is smaller than src" % (dest, src))
+
+    src_size = bits_required(src)
+
+    # Calculate the packed version
+    srcp = packing.pack(src & ((1<<32)-1), dest.size)
+
+    # Calculate the unsigned and signed versions
+    srcu = packing.unpack(srcp, dest.size, sign=False)
+    srcs = packing.unpack(srcp, dest.size, sign=True)
+
+%>\
+% if is_register(src):
+    % if src == dest:
+    /* moving ${src} into ${dest}, but this is a no-op */
+    % elif src.name in dest.bigger:
+    /* moving ${src} into ${dest}, but this is a no-op */
+    % elif dest.size > src.size:
+    movzx ${dest}, ${src}
+    % else:
+    mov ${dest}, ${src}
+    % endif
 % elif isinstance(src, (int, long)):
-    <%
-     if src >= 2**sizes[dest] or src < -(2**(sizes[dest]-1)):
-         log.error('Number 0x%x does not fit into %s' % (src, dest))
-         return
-
-     # Calculate the unsigned and signed versions
-     srcu = src & (2 ** sizes[dest] - 1)
-     srcs = srcu - 2 * (srcu & (2 ** (sizes[dest] - 1)))
-
-     # Calculate the packed version
-     srcp = packing.pack(srcu, sizes[dest], 'little', False)
-    %>\
+## Special case for zeroes
     % if src == 0:
         xor ${dest}, ${dest}
-    % elif src == 10 and stack_allowed and sizes[dest] == 32:
-        push 9
-        pop ${bigger[dest][0]}
-        inc ${dest}
-    % elif stack_allowed and sizes[dest] == 32 and -128 <= srcs <= 127 and okay(srcp[0]):
-        push ${pretty(srcs)}
+## Special case for *just* a newline
+    % elif stack_allowed and dest.size == 32 and src == 10:
+        push 9 /* mov ${dest}, '\n' */
         pop ${dest}
+        inc ${dest}
+## Can we push/pop it?
+## This is shorter than a `mov` and has a better (more ASCII) encoding.
+## Note there are two variants for PUSH IMM32 and PUSH IMM8
+    % elif stack_allowed and dest.size == 32 and okay(srcp):
+        push ${pretty(src)}
+        pop ${dest}
+    % elif stack_allowed and dest.size == 32 and  -127 <= srcs < 128 and okay(srcp[0]):
+        push ${pretty(src)}
+        pop ${dest}
+## Easy case, everybody is happy
     % elif okay(srcp):
         mov ${dest}, ${pretty(src)}
-    % elif 0 <= srcu < 2**8 and okay(srcp[0]) and sizes[smaller[dest][-1]] == 8:
+## If it's an IMM8, we can use the 8-bit register
+    % elif 0 <= srcu < 2**8 and okay(srcp[0]) and dest.sizes[8]:
         xor ${dest}, ${dest}
-        mov ${smaller[dest][-1]}, ${pretty(srcu)}
-    % elif srcu == srcu & 0xff00 and okay(srcp[1]) and sizes[smaller[dest][-2]] == 8:
+        mov ${dest.sizes[8]}, ${pretty(srcu)}
+## If it's an IMM16, but there's nothing in the lower 8 bits,
+## we can use the high-8-bits register.
+## However, we must check that it exists.
+    % elif srcu == (srcu & 0xff00) and okay(srcp[1]) and dest.ff00:
+        xor ${dest}, ${dest} /* mov ${dest}, ${pretty(src)} */
+        mov ${dest.ff00}, ${pretty(srcu >> 8)}
+## If it's an IMM16, use the 16-bit register
+    % elif 0 <= srcu < 2**16 and okay(srcp[:2]) and dest.sizes[16]:
         xor ${dest}, ${dest}
-        mov ${smaller[dest][-2]}, ${pretty(srcu >> 8)}
-    % elif 0 <= srcu < 2**16 and okay(srcp[:2]) and sizes[smaller[dest][0]] == 16:
-        xor ${dest}, ${dest}
-        mov ${smaller[dest][0]}, ${pretty(src)}
+        mov ${dest.sizes[16]}, ${pretty(src)}
+## We couldn't find a way to make things work out, so just do
+## the XOR trick.
     % else:
         <%
         a,b = fiddling.xor_pair(srcp, avoid = '\x00\n')
-        a = hex(packing.unpack(a, sizes[dest]))
-        b = hex(packing.unpack(b, sizes[dest]))
+        a = hex(packing.unpack(a, dest.size))
+        b = hex(packing.unpack(b, dest.size))
         %>\
-        mov ${dest}, ${a}
+        mov ${dest}, ${a} /* mov ${dest}, ${src_name} */
         xor ${dest}, ${b}
     % endif
-% elif src in all_regs:
-    % if src == dest or src in bigger[dest] or src in smaller[dest]:
-        /* moving ${src} into ${dest}, but this is a no-op */
-    % elif sizes[dest] == sizes[src]:
-        mov ${dest}, ${src}
-    % elif sizes[dest] > sizes[src]:
-        movzx ${dest}, ${src}
-    % else:
-        <% done = False %>\
-        % for r in reversed(smaller[src]):
-            % if sizes[r] == sizes[dest]:
-                mov ${dest}, ${r}
-                <% done = True %>\
-                <% break %>\
-            % endif
-        % endfor
-        % if not done:
-            <% log.error('Register %s could not be moved into %s' % (src, dest)) %>\
-        % endif
-    % endif
-% else:
-    <% log.error('%s is neither a register nor an immediate' % src) %>\
 % endif