1212
1313import lief #type:ignore
1414
15+ # temporary constant, to be replaced with lief.ELF.ARCH.RISCV
16+ # https://github.com/lief-project/LIEF/pull/562
17+ LIEF_ELF_ARCH_RISCV = lief .ELF .ARCH (243 )
18+
1519def check_ELF_RELRO (binary ) -> bool :
1620 '''
1721 Check for read-only relocations.
@@ -178,31 +182,46 @@ def check_control_flow(binary) -> bool:
178182 return True
179183 return False
180184
181-
182- CHECKS = {
183- lief .EXE_FORMATS .ELF : [
185+ BASE_ELF = [
184186 ('PIE' , check_PIE ),
185187 ('NX' , check_NX ),
186188 ('RELRO' , check_ELF_RELRO ),
187189 ('Canary' , check_ELF_Canary ),
188190 ('separate_code' , check_ELF_separate_code ),
189- ],
190- lief .EXE_FORMATS .PE : [
191+ ]
192+
193+ BASE_PE = [
191194 ('PIE' , check_PIE ),
192195 ('DYNAMIC_BASE' , check_PE_DYNAMIC_BASE ),
193196 ('HIGH_ENTROPY_VA' , check_PE_HIGH_ENTROPY_VA ),
194197 ('NX' , check_NX ),
195198 ('RELOC_SECTION' , check_PE_RELOC_SECTION ),
196199 ('CONTROL_FLOW' , check_PE_control_flow ),
197- ],
198- lief .EXE_FORMATS .MACHO : [
200+ ]
201+
202+ BASE_MACHO = [
199203 ('PIE' , check_PIE ),
200204 ('NOUNDEFS' , check_MACHO_NOUNDEFS ),
201205 ('NX' , check_NX ),
202206 ('LAZY_BINDINGS' , check_MACHO_LAZY_BINDINGS ),
203207 ('Canary' , check_MACHO_Canary ),
204208 ('CONTROL_FLOW' , check_control_flow ),
205209]
210+
211+ CHECKS = {
212+ lief .EXE_FORMATS .ELF : {
213+ lief .ARCHITECTURES .X86 : BASE_ELF ,
214+ lief .ARCHITECTURES .ARM : BASE_ELF ,
215+ lief .ARCHITECTURES .ARM64 : BASE_ELF ,
216+ lief .ARCHITECTURES .PPC : BASE_ELF ,
217+ LIEF_ELF_ARCH_RISCV : BASE_ELF ,
218+ },
219+ lief .EXE_FORMATS .PE : {
220+ lief .ARCHITECTURES .X86 : BASE_PE ,
221+ },
222+ lief .EXE_FORMATS .MACHO : {
223+ lief .ARCHITECTURES .X86 : BASE_MACHO ,
224+ }
206225}
207226
208227if __name__ == '__main__' :
@@ -211,13 +230,24 @@ def check_control_flow(binary) -> bool:
211230 try :
212231 binary = lief .parse (filename )
213232 etype = binary .format
233+ arch = binary .abstract .header .architecture
234+ binary .concrete
235+
214236 if etype == lief .EXE_FORMATS .UNKNOWN :
215237 print (f'{ filename } : unknown executable format' )
216238 retval = 1
217239 continue
218240
241+ if arch == lief .ARCHITECTURES .NONE :
242+ if binary .header .machine_type == LIEF_ELF_ARCH_RISCV :
243+ arch = LIEF_ELF_ARCH_RISCV
244+ else :
245+ print (f'{ filename } : unknown architecture' )
246+ retval = 1
247+ continue
248+
219249 failed : List [str ] = []
220- for (name , func ) in CHECKS [etype ]:
250+ for (name , func ) in CHECKS [etype ][ arch ] :
221251 if not func (binary ):
222252 failed .append (name )
223253 if failed :
0 commit comments