1 files changed, 57 insertions, 11 deletions

diff --git a/contrib/devtools/security-check.py b/contrib/devtools/security-check.py
index ef421aebb1..137fe377da 100755
--- a/contrib/devtools/security-check.py
+++ b/contrib/devtools/security-check.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-# Copyright (c) 2015-2020 The Bitcoin Core developers
+# Copyright (c) 2015-2021 The Bitcoin Core developers
 # Distributed under the MIT software license, see the accompanying
 # file COPYING or http://www.opensource.org/licenses/mit-license.php.
 '''
@@ -12,6 +12,10 @@ from typing import List
 
 import lief #type:ignore
 
+# temporary constant, to be replaced with lief.ELF.ARCH.RISCV
+# https://github.com/lief-project/LIEF/pull/562
+LIEF_ELF_ARCH_RISCV = lief.ELF.ARCH(243)
+
 def check_ELF_RELRO(binary) -> bool:
     '''
     Check for read-only relocations.
@@ -121,6 +125,21 @@ def check_PE_RELOC_SECTION(binary) -> bool:
     '''Check for a reloc section. This is required for functional ASLR.'''
     return binary.has_relocations
 
+def check_PE_control_flow(binary) -> bool:
+    '''
+    Check for control flow instrumentation
+    '''
+    main = binary.get_symbol('main').value
+
+    section_addr = binary.section_from_rva(main).virtual_address
+    virtual_address = binary.optional_header.imagebase + section_addr + main
+
+    content = binary.get_content_from_virtual_address(virtual_address, 4, lief.Binary.VA_TYPES.VA)
+
+    if content == [243, 15, 30, 250]: # endbr64
+        return True
+    return False
+
 def check_MACHO_NOUNDEFS(binary) -> bool:
     '''
     Check for no undefined references.
@@ -163,23 +182,24 @@ def check_control_flow(binary) -> bool:
         return True
     return False
 
-
-CHECKS = {
-'ELF': [
+BASE_ELF = [
     ('PIE', check_PIE),
     ('NX', check_NX),
     ('RELRO', check_ELF_RELRO),
     ('Canary', check_ELF_Canary),
     ('separate_code', check_ELF_separate_code),
-],
-'PE': [
+]
+
+BASE_PE = [
     ('PIE', check_PIE),
     ('DYNAMIC_BASE', check_PE_DYNAMIC_BASE),
     ('HIGH_ENTROPY_VA', check_PE_HIGH_ENTROPY_VA),
     ('NX', check_NX),
-    ('RELOC_SECTION', check_PE_RELOC_SECTION)
-],
-'MACHO': [
+    ('RELOC_SECTION', check_PE_RELOC_SECTION),
+    ('CONTROL_FLOW', check_PE_control_flow),
+]
+
+BASE_MACHO = [
     ('PIE', check_PIE),
     ('NOUNDEFS', check_MACHO_NOUNDEFS),
     ('NX', check_NX),
@@ -187,6 +207,21 @@ CHECKS = {
     ('Canary', check_MACHO_Canary),
     ('CONTROL_FLOW', check_control_flow),
 ]
+
+CHECKS = {
+    lief.EXE_FORMATS.ELF: {
+        lief.ARCHITECTURES.X86: BASE_ELF,
+        lief.ARCHITECTURES.ARM: BASE_ELF,
+        lief.ARCHITECTURES.ARM64: BASE_ELF,
+        lief.ARCHITECTURES.PPC: BASE_ELF,
+        LIEF_ELF_ARCH_RISCV: BASE_ELF,
+    },
+    lief.EXE_FORMATS.PE: {
+        lief.ARCHITECTURES.X86: BASE_PE,
+    },
+    lief.EXE_FORMATS.MACHO: {
+        lief.ARCHITECTURES.X86: BASE_MACHO,
+    }
 }
 
 if __name__ == '__main__':
@@ -194,14 +229,25 @@ if __name__ == '__main__':
     for filename in sys.argv[1:]:
         try:
             binary = lief.parse(filename)
-            etype = binary.format.name
+            etype = binary.format
+            arch = binary.abstract.header.architecture
+            binary.concrete
+
             if etype == lief.EXE_FORMATS.UNKNOWN:
                 print(f'{filename}: unknown executable format')
                 retval = 1
                 continue
 
+            if arch == lief.ARCHITECTURES.NONE:
+                if binary.header.machine_type == LIEF_ELF_ARCH_RISCV:
+                    arch = LIEF_ELF_ARCH_RISCV
+                else:
+                    print(f'{filename}: unknown architecture')
+                    retval = 1
+                    continue
+
             failed: List[str] = []
-            for (name, func) in CHECKS[etype]:
+            for (name, func) in CHECKS[etype][arch]:
                 if not func(binary):
                     failed.append(name)
             if failed: