Refactor encoders

pwnlib/encoders/arm/alphanumeric/ARM_Instructions.py

@@ -1,4 +1,5 @@
 from __future__ import division
+import six
 
 # +------------------------------------------------------------------------+ 
 # |                        ARM Instructions                                | 
@@ -18,122 +19,85 @@
 LSR = 11
 
 # (EOR/SUB/RSB)(PL/MI){S} rd, rn, #imm 
-# ==================================== 
+# ====================================
 def dpimm(op, cond, s, d, n, imm):
-   if type(imm) == int:
-      x = chr(imm & 0xff)
-   else:
-      x = imm
-   x += chr((d << 4) & 0xff)
-   if s:
-      if op == EOR:
-         x += chr(0x30 | n)
-      if op == SUB:
-         x += chr(0x50 | n)
-      if op == RSB:
-         x += chr(0x70 | n)
-   else:
-      if op == SUB:
-         x += chr(0x40 | n)
-      if op == RSB:
-         x += chr(0x60 | n)
-   if cond == PL:
-      x += "\x52"
-   else:
-      x += "\x42"
-   return x
-
+    x = bytearray()
+    if isinstance(imm, six.integer_types):
+        x.append(imm & 0xff)
+    else:
+        x.append(imm)
+    x.append((d << 4) & 0xff)
+    x.append(op << 5 | s << 4 | n)
+    x.append(cond << 4 | 2)
+    return bytes(x)
+
 # (EOR/SUB/RSB)PL{S} rd, rn, ra ROR #imm 
-# ====================================== 
+# ======================================
 def dpshiftimm(op, s, d, n, a, imm):
-   x = chr(0x60 | a)
-   x += chr(((d << 4)| (imm >> 1)) & 0xff)
-   if s:
-      if op == EOR:
-         x += chr(0x30 | n)
-      if op == SUB:
-         x += chr(0x50 | n)
-      if op == RSB:
-         x += chr(0x70 | n)
-   else:
-      if op == SUB:
-         x += chr(0x40 | n)
-      if op == RSB:
-         x += chr(0x60 | n)
-   return x + "\x50" 
+    x = bytearray()
+    x.append(0x60 | a)
+    x.append(((d << 4)| (imm >> 1)) & 0xff)
+    x.append(op << 5 | s << 4 | n)
+    x.append(PL << 4)
+    return bytes(x)
 
 # (EOR/SUB/RSB)PL{S} rd, rn, ra (ROR/LSR) rb 
-# ========================================== 
+# ==========================================
 def dpshiftreg(op, s, d, n, a, shift, b):
-   x = ''
-   if shift == LSR:
-      x += chr(0x30 | a)
-   else:
-      x += chr(0x70 | a)
-   x += chr(((d << 4) | b) & 0xff)
-   if s != 0:
-      if op == EOR:
-         x += chr(0x30 | n)
-      if op == SUB:
-         x += chr(0x50 | n)
-      if op == RSB:
-         x += chr(0x70 | n)
-   else:
-      if op == SUB:
-         x += chr(0x40 | n)
-      if op == RSB:
-         x += chr(0x60 | n)
-   return x + "\x50"
+    x = bytearray()
+    if shift == LSR:
+        x.append(0x30 | a)
+    else:
+        x.append(0x70 | a)
+    x.append(((d << 4) | b) & 0xff)
+    x.append(op << 5 | s << 4 | n)
+    x.append(PL << 4)
+    return bytes(x)
 
 # (LDR/STR)(PL/MI)B rd, [rn, #-imm] 
-# ================================= 
+# =================================
 def lsbyte(op, cond, d, n, imm):
-   if type(imm) == int:
-      x = chr(imm & 0xff)
-   else:
-      x = imm
-   x += chr((d << 4) & 0xff)
-#   x = chr(imm) + chr((d << 4) & 0xff)
-   if op == STR:
-      x += chr(0x40 | n)
-   else:
-      x += chr(0x50 | n)
-   if cond == PL:
-      x += "\x55"
-   else:
-      x += "\x45"
-   return x
+    x = bytearray()
+    if isinstance(imm, six.integer_types):
+        x.append(imm & 0xff)
+    else:
+        x.append(imm)
+    x.append((d << 4) & 0xff)
+    if op == STR:
+        x.append(0x40 | n)
+    else:
+        x.append(0x50 | n)
+    x.append(cond << 4 | 5)
+    return bytes(x)
 
 # STMPLFD rd, (Register List)^ 
-# ============================ 
+# ============================
 def smul(d, reglH, reglL):
-   return chr(reglL) + chr(reglH) + chr(0x40 | d) + "\x59"
+    return bytes(bytearray((reglL, reglH, 0x40 | d, 0x59)))
 
 # LDMPLDB rn!, (Register List) 
-# ============================ 
+# ============================
 def lmul(n, reglH, reglL):
-   return chr(reglL) + chr(reglH) + chr(0x30 | n) + "\x59"
+    return bytes(bytearray((reglL, reglH, 0x30 | n, 0x59)))
 
 # SWI(PL/MI) 0x9f0002 
-# ============== 
+# ==============
 def swi(cond):
-   x = "\x02\x00\x9f"
-   if cond == MI:
-      x += "\x4f"
-   else:
-      x += "\x5f"
-   return x
+    x = bytearray(b"\x02\x00\x9f")
+    x.append(cond << 4 | 0xf)
+    return bytes(x)
 
 # BMI 0xfffff4 
-# ============ 
+# ============
 def bmi():
-   return "\xf4\xff\xff\x4b"
+    return b"\xf4\xff\xff\x4b"
 
 # STRPLB rd, [!rn, -(rm ROR #imm)] with P=0 i.e. post-indexed addressing mode 
-# =========================================================================== 
+# ===========================================================================
 def sbyteposti(d, n, m, imm):
-   x = chr(0x60 | m)
-   x += chr(((d << 4) | (imm >> 1)) & 0xff)
-   x += chr(0x40 | n)
-   x += "\x56"
-   return x
+    x = bytearray()
+    x.append(0x60 | m)
+    x.append(((d << 4) | (imm >> 1)) & 0xff)
+    x.append(0x40 | n)
+    x.append(PL << 4 | 6)
+    return bytes(x)

pwnlib/encoders/arm/alphanumeric/__init__.py

@@ -14,8 +14,8 @@
 class ArmEncoder(Encoder):
     arch = 'arm'
 
-    blacklist  = {chr(c) for c in range(256) if chr(c) in (string.ascii_letters + string.digits)}
-    icache_flush = 1
+    blacklist  = set((string.ascii_letters + string.digits).encode())
+    icache_flush = True
 
     def __call__(self, input, avoid, pcreg=None):
         # If randomization is disabled, ensure that the seed
@@ -41,7 +41,7 @@ def __call__(self, input, avoid, pcreg=None):
         finally:
             random.setstate(state)
 
-        return output.encode()
+        return output
 
 class ThumbEncoder(ArmEncoder):
     arch = 'thumb'

pwnlib/encoders/arm/alphanumeric/alphanum_byte.py

@@ -1,52 +1,53 @@
 from __future__ import division
 from __future__ import absolute_import
 
+import six
 from . import random_funcs
 
 # +------------------------------------------------------------------------+ 
 # |                    ALPHANUMERIC MANIPULATIONS FUNCTIONS                | 
 # +------------------------------------------------------------------------+ 
 
-ALPHANUMERIC_BYTES = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+ALPHANUMERIC_BYTES = bytearray(b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz")
 
 # return 1 if the byte is alphanumeric 
 # ==================================== 
 def alphanumeric_check(c):
-   if type(c) == int:
-      c = chr(c & 0xff)
-   return c.isalnum()
+    if isinstance(c, six.integer_types):
+        return c&0xff in ALPHANUMERIC_BYTES
+    return c.isalnum()
 
 
 # return a random alphanumeric byte 
 # ================================= 
 def alphanumeric_get_byte():
-   return ord(random_funcs.randel(ALPHANUMERIC_BYTES))
+    return random_funcs.randel(ALPHANUMERIC_BYTES)
 
-# return a randomly selected alphanumeric byte less than max 
-# ========================================================== 
-#CSE author actually returns a byte <= max, not strictly < max
+# return a randomly selected alphanumeric byte less than max
+# ==========================================================
+# CSE author actually returns a byte <= max, not strictly < max
 def alphanumeric_get_byte_ltmax(max):
-   sz = 0
-   while sz < len(ALPHANUMERIC_BYTES) and ord(ALPHANUMERIC_BYTES[sz]) <= max:
-      sz += 1
-   return ord(random_funcs.randel(ALPHANUMERIC_BYTES[:sz]))
+    sz = 0
+    while sz < len(ALPHANUMERIC_BYTES) and ALPHANUMERIC_BYTES[sz] <= max:
+        sz += 1
+    return random_funcs.randel(ALPHANUMERIC_BYTES[:sz])
 
-# generate an alphanumeric offset such that c+offset is also alphanumeric 
-# ======================================================================= 
+# generate an alphanumeric offset such that c+offset is also alphanumeric
+# =======================================================================
 def off_gen(c):
-   if c >= 0 and c <= 0x4a:
-      max = 16 * 7 + 10 - c
-      while True:
-         x = alphanumeric_get_byte_ltmax(max)
-         if alphanumeric_check(c + x):
-            return x
-   return 0
+    if c >= 0 and c <= 0x4a:
+        max = 16 * 7 + 10 - c
+        while True:
+            x = alphanumeric_get_byte_ltmax(max)
+            if alphanumeric_check(c + x):
+                return x
+    return 0
 
 # return an alphanumeric value ret such that c XOR ret is also alphanumeric
 # =========================================================================
 def alphanumeric_get_complement(c):
-   c &= 0xff
-   while True:
-      ret = alphanumeric_get_byte()
-      if alphanumeric_check(c ^ ret):
-         return ret
+    c &= 0xff
+    while True:
+        ret = alphanumeric_get_byte()
+        if alphanumeric_check(c ^ ret):
+            return ret