Implement ELF and assembler

ptrlib/arch/arm/__init__.py

@@ -0,0 +1,3 @@
+from .archname import *
+from .assembler import *
+from .disassembler import *

ptrlib/arch/arm/archname.py

@@ -0,0 +1,30 @@
+from logging import getLogger
+
+logger = getLogger(__name__)
+
+
+def is_arch_arm(arch, bits=None):
+    """Check if architecture name string is ARM series
+
+    Args:
+        arch (str): Architecture name
+        bits (int): 32 or 64 (None by default)
+
+    Returns:
+        tuple: Returns tuple of canonicalized bits and name, or None.
+    """
+    arch = arch.lower().replace(' ', '').replace('_', '-')
+
+    if bits is not None and bits != 16 and bits != 32 and bits != 64:
+        logger.warn(f"Unknown bits: expected 16/32/64 but {bits} is given")
+        raise ValueError("Unknown architecture '{}:{}'".format(arch, bits))
+
+    if arch in ('arm', 'arm32', 'aarch32'):
+        # ARM
+        return True
+
+    elif arch in ['aarch', 'arm64', 'aarch64']:
+        # AArch64
+        return True
+
+    return False

ptrlib/arch/arm/assembler.py

@@ -0,0 +1,68 @@
+import contextlib
+import os
+import platform
+import subprocess
+import tempfile
+from logging import getLogger
+
+logger = getLogger(__name__)
+
+
+def assemble_arm(code, bits, entry, gcc_path=None, objcopy_path=None):
+    """Assemble code to intel machine code
+
+    Args:
+       code (bytes): Assembly code
+       bits (int): Bits of architecture
+       entry (bytes): Entry point
+    """
+    from ptrlib.arch.common import which
+    from .archname import is_arch_arm
+
+    if gcc_path is None or objcopy_path is None:
+        if is_arch_arm(platform.machine()):
+            # arm --> arm: Use native compiler
+            # TODO: Handle 32/64 bits difference
+            logger.warn("This feature is not fully implemented")
+            gcc_path = which('gcc')
+            objcopy_path = which('objcopy')
+        else:
+            # not-arm --> arm: Use corss-platform compiler
+            if bits == 32:
+                gcc_path = which('arm-linux-gnueabi-gcc')
+                objcopy_path = which('arm-linux-gnueabi-objcopy')
+            else:
+                gcc_path = which('aarch64-linux-gnu-gcc')
+                objcopy_path = which('aarch64-linux-gnu-objcopy')
+
+    fname_s   = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())+'.S'
+    fname_o   = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())+'.o'
+    fname_bin = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())+'.bin'
+    with open(fname_s, 'wb') as f:
+        f.write(code)
+
+    with contextlib.suppress(FileNotFoundError):
+        # Assemble
+        cmd = [gcc_path, '-nostdlib', '-c', fname_s, '-o', fname_o]
+        cmd.append('-Wl,--entry={}'.format(entry))
+        if subprocess.Popen(cmd).wait() != 0:
+            logger.warn("Assemble failed")
+            os.unlink(fname_s)
+            return
+
+        # Extract
+        cmd = [objcopy_path, '-O', 'binary', '-j', '.text', fname_o, fname_bin]
+        if subprocess.Popen(cmd).wait() != 0:
+            logger.warn("Extract failed")
+            os.unlink(fname_s)
+            os.unlink(fname_o)
+            return
+
+        with open(fname_bin, 'rb') as f:
+            output = f.read()
+
+        os.unlink(fname_s)
+        os.unlink(fname_o)
+        os.unlink(fname_bin)
+
+        return output

ptrlib/arch/common/__init__.py

@@ -1 +1,3 @@
+from .assembler import *
+from .disassembler import *
 from .ospath import *

ptrlib/arch/common/assembler.py

@@ -0,0 +1,69 @@
+import contextlib
+import os
+import subprocess
+import tempfile
+from logging import getLogger
+from ptrlib.arch.intel import assemble_intel, is_arch_intel
+from ptrlib.arch.arm   import assemble_arm, is_arch_arm
+from ptrlib.binary.encoding import *
+
+logger = getLogger(__name__)
+
+
+def assemble(code, bits=None, arch='intel', syntax='intel', entry=None,
+             as_path=None, ld_path=None):
+    if isinstance(code, str):
+        code = str2bytes(code)
+
+    if code[-1] != 0x0a:
+        code += b'\n'
+
+    if entry is None:
+        entry = 'ptrlib_main'
+        code = b'.global ptrlib_main\nptrlib_main:\n' + code
+
+    if is_arch_intel(arch, bits):
+        if syntax.lower() == 'intel':
+            code = b'.intel_syntax noprefix\n' + code
+        return assemble_intel(code, bits, entry, as_path, ld_path)
+
+    elif is_arch_arm(arch, bits):
+        return assemble_arm(code, bits, entry, as_path, ld_path)
+
+    else:
+        raise ValueError("Unknown architecture '{}'".format(arch))
+
+
+def nasm(code, fmt='bin', bits=None, org=None, nasm_path=None):
+    from ptrlib.arch.common import which
+    if nasm_path is None:
+        nasm_path = which('nasm')
+
+    if isinstance(code, str):
+        code = str2bytes(code)
+
+    if bits is not None:
+        code = 'BITS {}\n'.format(bits).encode() + code
+    if org is not None:
+        code = 'ORG {}\n'.format(org).encode() + code
+
+    fname_s = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())+'.S'
+    fname_o = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())+'.o'
+    with open(fname_s, 'wb') as f:
+        f.write(code)
+
+    with open(fname_o, 'wb+') as f, contextlib.suppress(FileNotFoundError):
+        p = subprocess.Popen([nasm_path, "-f{}".format(fmt),
+                              fname_s, "-o", fname_o])
+        if p.wait() != 0:
+            logger.warn("Assemble failed")
+            os.unlink(fname_s)
+            return None
+
+        f.seek(0)
+        output = f.read()
+
+        os.unlink(fname_s)
+        os.unlink(fname_o)
+
+        return output

ptrlib/arch/common/disassembler.py

@@ -0,0 +1,20 @@
+import subprocess
+from logging import getLogger
+
+logger = getLogger(__name__)
+
+
+def disassemble(code, bits=None, arch='intel'):
+    raise NotImplementedError()
+
+def disasm(code,
+           arch='x86',
+           mode='64',
+           endian='little',
+           address=0,
+           micro=False,
+           mclass=False,
+           v8=False,
+           v9=False,
+           returns=list):
+    raise NotImplementedError()

ptrlib/arch/common/ospath.py

@@ -3,4 +3,4 @@
 
 def which(s):
     # TODO: Separate Windows support based on running OS
-    return linux_which(s)
+    return which_linux(s)

ptrlib/arch/intel/__init__.py

@@ -0,0 +1,3 @@
+from .archname import *
+from .assembler import *
+from .disassembler import *

ptrlib/arch/intel/archname.py

@@ -0,0 +1,30 @@
+from logging import getLogger
+
+logger = getLogger(__name__)
+
+
+def is_arch_intel(arch, bits=None):
+    """Check if architecture name string is intel series
+
+    Args:
+        arch (str): Architecture name
+        bits (int): 32 or 64 (None by default)
+
+    Returns:
+        tuple: Returns tuple of canonicalized bits and name, or None.
+    """
+    arch = arch.lower().replace(' ', '').replace('_', '-')
+
+    if bits is not None and bits != 16 and bits != 32 and bits != 64:
+        logger.warn(f"Unknown bits: expected 16/32/64 but {bits} is given")
+        raise ValueError("Unknown architecture '{}:{}'".format(arch, bits))
+
+    if arch in ('intel', 'intel32', 'i386', 'x86'):
+        # x86
+        return True
+
+    elif arch in ['intel64', 'x86-64', 'x64', 'amd', 'amd64']:
+        # x86-64
+        return True
+
+    return False

ptrlib/arch/intel/assembler.py

@@ -0,0 +1,65 @@
+import contextlib
+import os
+import platform
+import subprocess
+import tempfile
+from logging import getLogger
+
+logger = getLogger(__name__)
+
+
+def assemble_intel(code, bits, entry, gcc_path=None, objcopy_path=None):
+    """Assemble code to intel machine code
+
+    Args:
+       code (bytes): Assembly code
+       bits (int): Bits of architecture
+       entry (bytes): Entry point
+    """
+    from ptrlib.arch.common import which
+    from .archname import is_arch_intel
+
+    if gcc_path is None or objcopy_path is None:
+        if is_arch_intel(platform.machine()):
+            # intel --> intel: Use native compiler
+            gcc_path = which('gcc')
+            objcopy_path = which('objcopy')
+        else:
+            # not-intel --> intel: Use corss-platform compiler
+            gcc_path = which('x86_64-linux-gnu-gcc')
+            objcopy_path = which('x86_64-linux-gnu-objcopy')
+
+    if bits == 32:
+        code = b'.code32\n' + code
+
+    fname_s   = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())+'.S'
+    fname_o   = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())+'.o'
+    fname_bin = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())+'.bin'
+    with open(fname_s, 'wb') as f:
+        f.write(code)
+
+    with contextlib.suppress(FileNotFoundError):
+        # Assemble
+        cmd = [gcc_path, '-nostdlib', '-c', fname_s, '-o', fname_o]
+        cmd.append('-Wl,--entry={}'.format(entry))
+        if subprocess.Popen(cmd).wait() != 0:
+            logger.warn("Assemble failed")
+            os.unlink(fname_s)
+            return
+
+        # Extract
+        cmd = [objcopy_path, '-O', 'binary', '-j', '.text', fname_o, fname_bin]
+        if subprocess.Popen(cmd).wait() != 0:
+            logger.warn("Extract failed")
+            os.unlink(fname_s)
+            os.unlink(fname_o)
+            return
+
+        with open(fname_bin, 'rb') as f:
+            output = f.read()
+
+        os.unlink(fname_s)
+        os.unlink(fname_o)
+        os.unlink(fname_bin)
+
+        return output

ptrlib/arch/linux/ospath.py

@@ -1,7 +1,7 @@
 import subprocess
 
 
-def linux_which(s):
+def which_linux(s):
     if '/' not in s:
         try:
             s = subprocess.check_output(["which", s]).decode().rstrip()

ptrlib/crypto/blockcipher/__init__.py

@@ -1,2 +1,3 @@
 from .ecb import *
 from .padcbc import *
+from .padding import *

ptrlib/crypto/blockcipher/padding.py

@@ -0,0 +1,41 @@
+from logging import getLogger
+
+logger = getLogger(__name__)
+
+
+def pad(data, size, mode='pkcs#5'):
+    """Append Padding
+
+    Args:
+        data (bytes) : Data to append padding to
+        size (int)   : Block size
+        mode (str)   : Padding mode
+
+    Available modes:
+      zero  : Zero byte padding
+      pkcs#5: PKCS#5 Padding
+    """
+    mode = mode.lower()
+    if mode not in ['zero', 'pkcs#5', '']:
+        logger.warning("Invalid padding mode. Using 'zero'")
+        logger.warning("Choose from zero / pkcs#5 / ")
+        mode = 'zero'
+
+    if isinstance(data, str):
+        data = str2bytes(data)
+
+    if size <= 0:
+        logger.warning("Block size must be bigger than zero")
+        return data
+
+    if mode == 'zero':
+        # Zero byte padding
+        return data + b'\x00' * (size - (len(data) % size))
+
+    elif mode == 'pkcs#5':
+        # PKCS#5
+        padlen = size - (len(data) % size)
+        if padlen > 255:
+            logger.warning("Padding length cannot be bigger than 0xff in PKCS#5")
+            padlen %= 0x100
+        return data + bytes([padlen]) * padlen