pyasm.py

#!/usr/bin/env python3
import os, tempfile, shutil, subprocess, ctypes

class BaseAsmFunc(object):
    def __init__(self, sfunc):
        self.source = sfunc.__doc__
        self._tmp = tempfile.mkdtemp() + os.sep
        self.compile(self.source)

    def compile(self, source):
        self.sfile = self._tmp + "b.S"
        with open(self.sfile, "w") as fd:
            fd.write(self.HEADER + "\n")
            fd.write(source + "\n")
            fd.write(self.FOOTER + "\n")
        self.sofile = self._tmp + "b.so"
        subprocess.check_call("gcc %s -shared -o %s %s" % (self.CFLAGS, self.sofile, self.sfile), shell=True)
        self.lib = ctypes.CDLL(self.sofile)

    def __call__(self, *args):
        return self.lib._start(*args)

    def __del__(self):
        if self._tmp:
            shutil.rmtree(self._tmp)
            self._tmp = None

class X64AsmFunc(BaseAsmFunc):
    CFLAGS = "-pipe -Wall"
    HEADER = """
    .text
    .code64
    .globl _start
_start:
    """
    FOOTER = """
    """

if __name__ == "__main__":
    import struct

    @X64AsmFunc
    def cpuid_eax():
        """
        push %rbx
        mov %edi, %eax
        cpuid
        pop %rbx
        ret
        """

    @X64AsmFunc
    def cpuid_ebx():
        """
        push %rbx
        mov %edi, %eax
        cpuid
        mov %ebx, %eax
        pop %rbx
        ret
        """

    @X64AsmFunc
    def cpuid_ecx():
        """
        push %rbx
        mov %edi, %eax
        cpuid
        mov %ecx, %eax
        pop %rbx
        ret
        """

    @X64AsmFunc
    def cpuid_edx():
        """
        push %rbx
        mov %edi, %eax
        cpuid
        mov %edx, %eax
        pop %rbx
        ret
        """

    print("CPU vendor: %s" % struct.pack("<III", cpuid_ebx(0), cpuid_edx(0),
                                         cpuid_ecx(0)).decode("ascii"))