Add experimental support for consts

Author: ptr-yudai

README.md

@@ -22,11 +22,20 @@ Library Dependency:
 - pywin32 (when handling Windows process)
 
 External Program:
-- When using `SSH` function:
+- `SSH` requires:
   - ssh
   - expect
-- When using `nasm` function:
+- `nasm` requires:
   - nasm
+- `assemble` requires:
+  - gcc, objcopy (x86, x86-64)
+  - arm-linux-gnueabi-gcc, aarch64-linux-gnu-gcc (arm, aarch64)
+- `disassemble` requires:
+  - objdump (x86, x86-64)
+  - arm-linux-gnueabi-objdump, aarch64-linux-gnu-objdump (arm, aarch64)
+- `consts` requires:
+  - grep
+  - gcc (x86, x86-64)
 
 ## Usage
 Basic examples are available at [/examples](https://github.com/ptr-yudai/ptrlib/tree/master/examples/).

ptrlib/arch/arm/__init__.py

@@ -1,4 +1,5 @@
 from .archname import *
 from .assembler import *
+from .consts import *
 from .disassembler import *
 from .syscall import *

ptrlib/arch/arm/consts.py

@@ -0,0 +1,16 @@
+import functools
+from typing import Union
+
+try:
+    cache = functools.cache
+except AttributeError:
+    cache = functools.lru_cache
+
+
+class ConstsTableArm(object):
+    @cache
+    def __getitem__(self, name: str) -> Union[int, str]:
+        raise NotImplementedError("ARM is not supported yet")
+
+    def __getattr__(self, name: str):
+        return self[name]

ptrlib/arch/arm/syscall.py

@@ -1,5 +1,6 @@
 from .archname import *
 from .assembler import *
+from .consts import *
 from .cpu import *
 from .disassembler import *
 from .simd import *

ptrlib/arch/intel/__init__.py

@@ -0,0 +1,21 @@
+import functools
+from typing import Union
+
+try:
+    cache = functools.cache
+except AttributeError:
+    cache = functools.lru_cache
+
+
+class ConstsTableIntel(object):
+    @cache
+    def __getitem__(self, name: str) -> Union[int, str]:
+        from ptrlib.arch.linux import consts
+        return consts.resolve_constant(
+            name,
+            ['/usr/include/x86_64-linux-gnu',
+             '/usr/include/x86_64-linux-musl']
+        )
+
+    def __getattr__(self, name: str):
+        return self[name]

ptrlib/arch/intel/consts.py

@@ -1,3 +1,4 @@
+from .consts import *
 from .ospath import *
 from .sig import *
 from .syscall import *

ptrlib/arch/intel/syscall.py

@@ -0,0 +1,130 @@
+import contextlib
+import functools
+import os
+import platform
+import re
+import subprocess
+import tempfile
+from typing import List, Optional, Union
+from ptrlib.arch.arm import is_arch_arm, ConstsTableArm
+from ptrlib.arch.intel import is_arch_intel, ConstsTableIntel
+
+try:
+    cache = functools.cache
+except AttributeError:
+    cache = functools.lru_cache
+
+
+_TEMPLATE_C = """
+#include <stdio.h>
+#include <{0}>
+
+#define print_const(X) (void)_Generic((X),   \
+  char*: printf("S:%s\\n", (const char*)(X)), \
+  default: printf("V:%lu\\n", (size_t)(X))    \
+)
+
+int main() {{
+  print_const({1});
+  return 0;
+}}
+"""
+
+# ConstsTableLinux: Experimental feature
+class ConstsTableLinux(object):
+    def resolve_constant(self,
+                         const: str,
+                         include_path: Optional[List[str]] = None) -> Union[int, str]:
+        from ptrlib.arch.common import which
+
+        if len(const) == 0:
+            raise KeyError("Empty name '{}'".format(const))
+
+        if include_path is not None:
+            include_path = include_path + ['/usr/include']
+        else:
+            include_path = ['/usr/include']
+
+        def heuristic_redirect(path: str) -> str:
+            """Convert include path"""
+            with open(path, 'r') as f:
+                buf = f.read()
+                found = re.findall(r"Never use <.+> directly; include <(.+)> instead\.", buf)
+                if found:
+                    return found[0]
+                else:
+                    return path
+
+        def test_constant(path: str, name: str, gcc_path: str) -> Optional[Union[int, str]]:
+            """Compile and run C code to get constant value"""
+            path = heuristic_redirect(path)
+            fname_c   = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())+'.c'
+            fname_bin = os.path.join(tempfile.gettempdir(), os.urandom(24).hex())+'.bin'
+            with open(fname_c, 'w') as f:
+                f.write(_TEMPLATE_C.format(path, name))
+
+            with contextlib.suppress(FileNotFoundError):
+                p = subprocess.run([gcc_path, fname_c, '-o', fname_bin],
+                                   stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
+                os.unlink(fname_c)
+
+                if p.returncode == 0:
+                    p = subprocess.run([fname_bin],
+                                       stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
+                    os.unlink(fname_bin)
+
+                    if p.returncode == 0:
+                        if p.stdout.startswith(b"S:"):
+                            return p.stdout[2:].decode().strip()
+                        elif p.stdout.startswith(b"V:"):
+                            return int(p.stdout[2:])
+                        else:
+                            raise RuntimeError(f"Unexpected output: {p.stdout.decode()}")
+
+                        return
+
+        # We rely on grep since it's much faster
+        grep_path = which('grep')
+        if grep_path is None:
+            raise FileNotFoundError("'grep' not found")
+
+        if is_arch_intel(platform.machine()):
+            gcc_path = which('gcc')
+        else:
+            gcc_path = which('x86_64-linux-gnu-gcc')
+        if gcc_path is None:
+            raise FileNotFoundError("Install 'gcc' for this architecture")
+
+        for dpath in include_path:
+            # We can directly build regex since `const` is a valid Python variable name
+            p = subprocess.run([grep_path, '-E', f'#\\s*define\\s+{const}', '-rl', dpath],
+                               stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
+            if p.returncode != 0:
+                continue
+
+            for path in p.stdout.decode().split('\n'):
+                if not os.path.exists(path):
+                    continue
+
+                c = test_constant(path, const, gcc_path)
+                if c is not None:
+                    return c
+
+        raise KeyError("Could not find constant: {}".format(const))
+
+    @cache
+    def __getitem__(self, const_or_arch: str) -> Union[int, str, ConstsTableIntel, ConstsTableArm]:
+        if is_arch_intel(const_or_arch):
+            return ConstsTableIntel()
+        elif is_arch_arm(const_or_arch):
+            return ConstsTableArm()
+        elif const_or_arch.isupper():
+            return self.resolve_constant(const_or_arch)
+        else:
+            raise KeyError("Invalid name '{}'".format(arch))
+
+    def __getattr__(self, arch: str):
+        return self[arch]
+
+
+consts = ConstsTableLinux()

ptrlib/arch/linux/__init__.py

@@ -0,0 +1,24 @@
+import unittest
+from ptrlib import consts
+from logging import getLogger, FATAL
+
+
+class TestConsts(unittest.TestCase):
+    def setUp(self):
+        getLogger("ptrlib").setLevel(FATAL)
+
+    def test_consts(self):
+        self.assertEqual(consts['x86']['EFAULT'], 14)
+        self.assertEqual(consts['i386']['EPERM'], 1)
+        self.assertEqual(consts['O_RDWR'], 2)
+        self.assertEqual(consts['ELF_NOTE_SOLARIS'], "SUNW Solaris")
+
+        self.assertEqual(consts.i386.READ_IMPLIES_EXEC, 0x400000)
+        self.assertEqual(consts.x64.AT_RECURSIVE, 0x8000)
+        self.assertEqual(consts.amd64.ENOENT, 2)
+        self.assertEqual(consts.MAP_PRIVATE | consts.MAP_ANONYMOUS, 0x22)
+
+        with self.assertRaises(KeyError):
+            _ = consts.PTRLIB
+        with self.assertRaises(KeyError):
+            _ = consts.x64.PTRLIB

ptrlib/arch/linux/consts.py

@@ -5,15 +5,15 @@
 def syscall_table(arch, bits):
     table = {}
     r = requests.get(URL.format(arch=arch, bits=bits))
-    for num, _, name in re.findall("(\d+).+(common|64|i386).+\ssys_([a-z0-9_]+)", r.text):
+    for num, _, name in re.findall(r"(\d+).+(common|64|i386).+\ssys_([a-z0-9_]+)", r.text):
         table[name] = int(num)
     return table
 
 def syscall_table_arm64():
     # WTF
     table = {}
     r = requests.get(URL)
-    for name, num in re.findall("#define\s+__NR_([a-z0-9_]+)\s+(\d+)", r.text):
+    for name, num in re.findall(r"#define\s+__NR_([a-z0-9_]+)\s+(\d+)", r.text):
         table[name] = int(num)
     return table