Refactor pwnlib.encoders.encode() to be more nice to bytes objects

docs/source/encoders.rst

@@ -5,25 +5,8 @@
 :mod:`pwnlib.encoders` --- Encoding Shellcode
 ===============================================
 
-.. automodule:: pwnlib.encoders.encoder
-   :members:
-
-.. automodule:: pwnlib.encoders.i386.ascii_shellcode
-   :members:
-   :special-members:
-   :exclude-members: __init__
-
-.. automodule:: pwnlib.encoders.i386.xor
-   :members:
-
-.. automodule:: pwnlib.encoders.i386.delta
-   :members:
-
-.. automodule:: pwnlib.encoders.amd64.delta
+.. automodule:: pwnlib.encoders
    :members:
 
-.. automodule:: pwnlib.encoders.arm.xor
-   :members:
-
-.. automodule:: pwnlib.encoders.mips.xor
+.. automodule:: pwnlib.encoders.encoder
    :members:

examples/encoder.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+from pwn import *
+
+context.randomize = False
+context.log_level = 'error'
+
+def go():
+    info('========== %s ==========', context.arch)
+    with context.silent:
+        sc = asm(shellcraft.sh())
+        avoid=b'\x00\n\t '
+        enc = encode(sc, avoid=avoid, force=1)
+
+    assert not (byteset(avoid) & byteset(enc))
+    assert enc != sc
+
+    with context.silent:
+        io = ELF.from_bytes(enc).process()
+        io.sendline(b'whoami')
+
+    try:
+        info('%r', io.recvline() == b'pwntools\n')
+    except EOFError:
+        info('EOFError')
+
+    with context.silent:
+        io.close()
+
+
+context.clear(arch='i386')
+go()
+
+context.clear(arch='amd64')
+go()
+
+context.clear(arch='arm')
+go()
+
+context.clear(arch='mips')
+go()

pwnlib/encoders/__init__.py

@@ -1,14 +1,19 @@
 # -*- coding:utf-8 -*-
 """
 Encode shellcode to avoid input filtering and impress your friends!
+
+Note
+----
+
+Some of these methods will fail on various architectures.
+Your best bet is to use :func:`.encoder.encode` with an ``avoid=...``.
 """
 from __future__ import absolute_import
 
 from pwnlib.encoders import amd64
 from pwnlib.encoders import arm
 from pwnlib.encoders import i386
 from pwnlib.encoders import mips
-from pwnlib.encoders.encoder import Encoder
 from pwnlib.encoders.encoder import alphanumeric
 from pwnlib.encoders.encoder import encode
 from pwnlib.encoders.encoder import line

pwnlib/encoders/amd64/delta.py

@@ -2,7 +2,7 @@
 from __future__ import division
 
 from pwnlib.encoders.i386.delta import i386DeltaEncoder
-
+from pwnlib.util.misc import byteset
 
 class amd64DeltaEncoder(i386DeltaEncoder):
     r"""
@@ -41,7 +41,7 @@ class amd64DeltaEncoder(i386DeltaEncoder):
 '''
     arch      = 'amd64'
     raw       = b'H\x8d5\xf9\xff\xff\xffH\x83\xc6\x1a\xfcH\x89\xf7\xac\x93\xac(\xd8\xaa\x80\xeb\xacu\xf5'
-    blacklist = set(raw)
+    blacklist = byteset(raw)
 
 encode = amd64DeltaEncoder()
 __all__ = ['encode']

pwnlib/encoders/arm/__init__.py

@@ -1,4 +1,4 @@
 from __future__ import absolute_import
 
-from pwnlib.encoders.arm import alphanumeric
+# from pwnlib.encoders.arm import alphanumeric
 from pwnlib.encoders.arm import xor

pwnlib/encoders/arm/alphanumeric/__init__.py

@@ -10,13 +10,13 @@
 
 from pwnlib.context import context
 from . import builder
-from pwnlib.encoders.encoder import Encoder
+from pwnlib.encoders.encoder_class import Encoder
 
 
 class ArmEncoder(Encoder):
     arch = 'arm'
 
-    blacklist  = {chr(c) for c in range(256) if chr(c) in (string.ascii_letters + string.digits)}
+    blacklist  = {bytes([c]) for c in range(256) if chr(c) in (string.ascii_letters + string.digits)}
     icache_flush = 1
 
     def __call__(self, input, avoid, pcreg=None):

pwnlib/encoders/arm/xor.py

@@ -4,9 +4,10 @@
 from pwnlib import shellcraft
 from pwnlib.asm import asm
 from pwnlib.context import context
-from pwnlib.encoders.encoder import Encoder
+from pwnlib.encoders.encoder_class import Encoder
 from pwnlib.util.fiddling import xor_key
 from pwnlib.util.lists import group
+from pwnlib.util.misc import byteset
 from pwnlib.util.packing import u8
 
 
@@ -45,7 +46,7 @@ class ArmXorEncoder(Encoder):
 payload:
     """
 
-    blacklist = set("\x01\x80\x03\x85\x04\x07\x87\x0c\x8f\x0f\x16\x1c\x9f\x84\xa0%$'-/\xb0\xbd\x81A@\xc2DG\xc6\xc8OPT\xd8_\xe1`\xe3\xe2\xe5\xe7\xe9\xe8\xea\xe0p\xf7")
+    blacklist = byteset(b"\x01\x80\x03\x85\x04\x07\x87\x0c\x8f\x0f\x16\x1c\x9f\x84\xa0%$'-/\xb0\xbd\x81A@\xc2DG\xc6\xc8OPT\xd8_\xe1`\xe3\xe2\xe5\xe7\xe9\xe8\xea\xe0p\xf7")
 
     def __call__(self, raw_bytes, avoid, pcreg=''):
         key, xordata = xor_key(raw_bytes, avoid, size=1)

pwnlib/encoders/encoder.py

@@ -2,49 +2,18 @@
 from __future__ import absolute_import
 from __future__ import division
 
-import collections
 import random
 import re
-import string
 
 from pwnlib.context import LocalContext
 from pwnlib.context import context
+from pwnlib.encoders.encoder_class import Encoder
 from pwnlib.log import getLogger
 from pwnlib.util.fiddling import hexdump
+from pwnlib.util.misc import byteset
 
 log = getLogger(__name__)
 
-class Encoder(object):
-    _encoders = collections.defaultdict(lambda: [])
-
-    #: Architecture which this encoder works on
-    arch = None
-
-    #: Blacklist of bytes which are known not to be supported
-    blacklist = set()
-
-    def __init__(self):
-        """Shellcode encoder class
-
-        Implements an architecture-specific shellcode encoder
-        """
-        Encoder._encoders[self.arch].append(self)
-
-    def __call__(self, raw_bytes, avoid, pcreg):
-        """avoid(raw_bytes, avoid)
-
-        Arguments:
-            raw_bytes(str):
-                String of bytes to encode
-            avoid(set):
-                Set of bytes to avoid
-            pcreg(str):
-                Register which contains the address of the shellcode.
-                May be necessary for some shellcode.
-        """
-        raise NotImplementedError()
-
-
 @LocalContext
 def encode(raw_bytes, avoid=None, expr=None, force=0, pcreg=''):
     """encode(raw_bytes, avoid, expr, force) -> str
@@ -54,37 +23,43 @@ def encode(raw_bytes, avoid=None, expr=None, force=0, pcreg=''):
 
     Arguments:
 
-        raw_bytes(str): Sequence of shellcode bytes to encode.
-        avoid(str):     Bytes to avoid
-        expr(str):      Regular expression which matches bad characters.
-        force(bool):    Force re-encoding of the shellcode, even if it
-                        doesn't contain any bytes in ``avoid``.
+        raw_bytes(bytes): Sequence of shellcode bytes to encode.
+        avoid(bytes):     Bytes to avoid
+        expr(bytes):      Regular expression which matches bad characters.
+        force(bool):      Force re-encoding of the shellcode, even if it
+                          doesn't contain any bytes in ``avoid``.
     """
     orig_avoid = avoid
 
-    avoid = set(avoid or '')
+    avoid = byteset(avoid or b'')
 
     if expr:
         for char in all_chars:
             if re.search(expr, char):
                 avoid.add(char)
 
-    if not (force or avoid & set(raw_bytes)):
+    if not (force or avoid & byteset(raw_bytes)):
         return raw_bytes
 
     encoders = Encoder._encoders[context.arch]
-    random.shuffle(encoders)
+
+    if context.randomize:
+        random.shuffle(encoders)
 
     for encoder in encoders:
         if encoder.blacklist & avoid:
             continue
 
+        log.debug('Selected encoder %r', encoder)
+
+        bytes_avoid = b''.join(avoid)
+
         try:
-            v = encoder(raw_bytes, bytes(avoid), pcreg)
+            v = encoder(raw_bytes, bytes_avoid, pcreg)
         except NotImplementedError:
             continue
 
-        if avoid & set(v):
+        if avoid & byteset(v):
             log.warning_once("Encoder %s did not succeed" % encoder)
             continue
 
@@ -97,23 +72,23 @@ def encode(raw_bytes, avoid=None, expr=None, force=0, pcreg=''):
     elif expr:
         avoid_errmsg = repr(expr)
     else:
-        avoid_errmsg = repr(bytes(avoid))
+        avoid_errmsg = repr(b''.join(avoid))
 
     args = (context.arch, avoid_errmsg, hexdump(raw_bytes))
     msg = "No encoders for %s which can avoid %s for\n%s" % args
     msg = msg.replace('%', '%%')
     log.error(msg)
 
-all_chars        = list(chr(i) for i in range(256))
-re_alphanumeric  = r'[^A-Za-z0-9]'
-re_printable     = r'[^\x21-\x7e]'
-re_whitespace    = r'\s'
-re_null          = r'\x00'
-re_line          = r'[\s\x00]'
+all_chars        = list(bytes([i]) for i in range(256))
+re_alphanumeric  = br'[^A-Za-z0-9]'
+re_printable     = br'[^\x21-\x7e]'
+re_whitespace    = br'\s'
+re_null          = br'\x00'
+re_line          = br'[\s\x00]'
 
 @LocalContext
 def null(raw_bytes, *a, **kw):
-    """null(raw_bytes) -> str
+    """null(raw_bytes) -> bytes
 
     Encode the shellcode ``raw_bytes`` such that it does not
     contain any NULL bytes.
@@ -124,7 +99,7 @@ def null(raw_bytes, *a, **kw):
 
 @LocalContext
 def line(raw_bytes, *a, **kw):
-    """line(raw_bytes) -> str
+    """line(raw_bytes) -> bytes
 
     Encode the shellcode ``raw_bytes`` such that it does not
     contain any NULL bytes or whitespace.
@@ -135,7 +110,7 @@ def line(raw_bytes, *a, **kw):
 
 @LocalContext
 def alphanumeric(raw_bytes, *a, **kw):
-    """alphanumeric(raw_bytes) -> str
+    """alphanumeric(raw_bytes) -> bytes
 
     Encode the shellcode ``raw_bytes`` such that it does not
     contain any bytes except for [A-Za-z0-9].
@@ -146,7 +121,7 @@ def alphanumeric(raw_bytes, *a, **kw):
 
 @LocalContext
 def printable(raw_bytes, *a, **kw):
-    """printable(raw_bytes) -> str
+    """printable(raw_bytes) -> bytes
 
     Encode the shellcode ``raw_bytes`` such that it only contains
     non-space printable bytes.
@@ -157,7 +132,7 @@ def printable(raw_bytes, *a, **kw):
 
 @LocalContext
 def scramble(raw_bytes, *a, **kw):
-    """scramble(raw_bytes) -> str
+    """scramble(raw_bytes) -> bytes
 
     Encodes the input data with a random encoder.
 

pwnlib/encoders/encoder_class.py

@@ -0,0 +1,37 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import
+from __future__ import division
+
+import collections
+
+from pwnlib.util.misc import byteset
+
+class Encoder(object):
+    _encoders = collections.defaultdict(lambda: [])
+
+    #: Architecture which this encoder works on
+    arch = None
+
+    #: Blacklist of bytes which are known not to be supported
+    blacklist = byteset()
+
+    def __init__(self):
+        """Shellcode encoder class
+
+        Implements an architecture-specific shellcode encoder
+        """
+        Encoder._encoders[self.arch].append(self)
+
+    def __call__(self, raw_bytes, avoid, pcreg):
+        """avoid(raw_bytes, avoid)
+
+        Arguments:
+            raw_bytes(bytes):
+                Bytes to encode
+            avoid(bytes):
+                Bytes to avoid
+            pcreg(str):
+                Register which contains the address of the shellcode.
+                May be necessary for some shellcode.
+        """
+        raise NotImplementedError()

pwnlib/encoders/i386/ascii_shellcode.py

@@ -10,9 +10,10 @@
 
 from pwnlib.context import LocalContext
 from pwnlib.context import context
-from pwnlib.encoders.encoder import Encoder
+from pwnlib.encoders.encoder_class import Encoder
 from pwnlib.encoders.encoder import all_chars
 from pwnlib.util.iters import group
+from pwnlib.util.misc import byteset
 from pwnlib.util.packing import *
 
 
@@ -94,13 +95,14 @@ def __call__(self, raw_bytes, avoid=None, pcreg=None):
             vocab = bytearray(
                 b"!\"#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz{|}~")
         else:
-            required_chars = set('\\-%TXP')
+            required_chars = byteset(b'\\-%TXP')
             allowed = set(all_chars)
+            avoid = byteset(avoid)
             if avoid.intersection(required_chars):
                 raise RuntimeError(
                     '''These characters ({}) are required because they assemble
                     into instructions used to unpack the shellcode'''.format(
-                        str(required_chars, 'ascii')))
+                        str(b''.join(required_chars), 'ascii')))
             allowed.difference_update(avoid)
             vocab = bytearray(map(ord, allowed))