mirror of
https://github.com/QuasarApp/LIEF.git
synced 2025-04-27 21:04:32 +00:00
154 lines
4.4 KiB
Python
154 lines
4.4 KiB
Python
#!/usr/bin/env python
|
|
import unittest
|
|
import logging
|
|
import os
|
|
import sys
|
|
import stat
|
|
import re
|
|
import subprocess
|
|
import tempfile
|
|
import shutil
|
|
from subprocess import Popen
|
|
|
|
import lief
|
|
from lief.ELF import Section
|
|
|
|
from unittest import TestCase
|
|
from utils import get_sample, get_compiler, is_aarch64, is_x86_64, is_linux
|
|
|
|
CURRENT_DIRECTORY = os.path.dirname(os.path.abspath(__file__))
|
|
STUB_FILE = None
|
|
if is_x86_64():
|
|
STUB_FILE = "hello_lief.bin"
|
|
elif is_aarch64():
|
|
STUB_FILE = "hello_lief_aarch64.bin"
|
|
|
|
STUB = lief.parse(os.path.join(CURRENT_DIRECTORY, STUB_FILE))
|
|
|
|
|
|
LIBADD_C = """\
|
|
#include <stdlib.h>
|
|
#include <stdio.h>
|
|
|
|
int add(int a, int b);
|
|
|
|
int add(int a, int b) {
|
|
printf("%d + %d = %d\\n", a, b, a+b);
|
|
return a + b;
|
|
}
|
|
"""
|
|
|
|
BINADD_C = """\
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
|
|
int add(int a, int b);
|
|
|
|
int main(int argc, char **argv) {
|
|
if (argc != 3) {
|
|
printf("Usage: %s <a> <b>\\n", argv[0]);
|
|
exit(-1);
|
|
}
|
|
|
|
int res = add(atoi(argv[1]), atoi(argv[2]));
|
|
printf("From myLIb, a + b = %d\\n", res);
|
|
return 0;
|
|
}
|
|
"""
|
|
|
|
|
|
class TestAddContent(TestCase):
|
|
def setUp(self):
|
|
self.logger = logging.getLogger(__name__)
|
|
self.tmp_dir = tempfile.mkdtemp(suffix='_lief_test_content')
|
|
self.logger.debug("temp dir: {}".format(self.tmp_dir))
|
|
|
|
self.binadd_path = os.path.join(self.tmp_dir, "binadd.c")
|
|
self.libadd_path = os.path.join(self.tmp_dir, "libadd.c")
|
|
|
|
self.libadd_so = os.path.join(self.tmp_dir, "libadd.so")
|
|
self.binadd_bin = os.path.join(self.tmp_dir, "binadd.bin")
|
|
|
|
self.compiler = get_compiler()
|
|
|
|
self.logger.debug("Compiler: {}".format(self.compiler))
|
|
|
|
with open(self.binadd_path, 'w') as f:
|
|
f.write(BINADD_C)
|
|
|
|
with open(self.libadd_path, 'w') as f:
|
|
f.write(LIBADD_C)
|
|
|
|
def compile_libadd(self, input, extra_flags=[]):
|
|
if os.path.isfile(self.libadd_so):
|
|
os.remove(self.libadd_so)
|
|
CC_FLAGS = ['-fPIC', '-shared'] + extra_flags
|
|
cmd = [self.compiler, '-o', self.libadd_so] + CC_FLAGS + [input]
|
|
self.logger.debug("Compile 'libadd' with: {}".format(" ".join(cmd)))
|
|
p = Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
|
|
stdout, _ = p.communicate()
|
|
self.logger.debug(stdout)
|
|
|
|
|
|
def compile_binadd(self, input, extra_flags=[]):
|
|
if os.path.isfile(self.binadd_bin):
|
|
os.remove(self.binadd_bin)
|
|
|
|
CC_FLAGS = ['-L', self.tmp_dir] + extra_flags
|
|
cmd = [self.compiler, '-o', self.binadd_bin] + CC_FLAGS + [input, '-ladd']
|
|
self.logger.debug("Compile 'binadd' with: {}".format(" ".join(cmd)))
|
|
p = Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
|
|
stdout, _ = p.communicate()
|
|
self.logger.debug(stdout)
|
|
|
|
|
|
@unittest.skipUnless(is_linux(), "requires Linux")
|
|
def test_simple(self):
|
|
self.compile_libadd(self.libadd_path)
|
|
self.compile_binadd(self.binadd_path)
|
|
|
|
libadd = lief.parse(self.libadd_so)
|
|
for i in range(10):
|
|
segment = libadd.add(STUB.segments[0])
|
|
segment.alignment = 0x1000
|
|
|
|
new_ep = (STUB.header.entrypoint - STUB.segments[0].virtual_address) + segment.virtual_address
|
|
|
|
if libadd.has(lief.ELF.DYNAMIC_TAGS.INIT_ARRAY):
|
|
init_array = libadd.get(lief.ELF.DYNAMIC_TAGS.INIT_ARRAY)
|
|
callbacks = init_array.array
|
|
callbacks[0] = new_ep
|
|
init_array.array = callbacks
|
|
|
|
if libadd.has(lief.ELF.DYNAMIC_TAGS.INIT):
|
|
init = libadd.get(lief.ELF.DYNAMIC_TAGS.INIT)
|
|
init.value = new_ep
|
|
|
|
libadd.write(self.libadd_so)
|
|
|
|
st = os.stat(self.libadd_so)
|
|
os.chmod(self.libadd_so, st.st_mode | stat.S_IEXEC)
|
|
|
|
p = Popen([self.binadd_bin, '1', '2'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env={"LD_LIBRARY_PATH": self.tmp_dir})
|
|
stdout, _ = p.communicate()
|
|
self.logger.debug(stdout.decode("utf8"))
|
|
|
|
self.assertIsNotNone(re.search(r'LIEF is Working', stdout.decode("utf8")))
|
|
|
|
def tearDown(self):
|
|
# Delete it
|
|
if os.path.isdir(self.tmp_dir):
|
|
shutil.rmtree(self.tmp_dir)
|
|
|
|
if __name__ == '__main__':
|
|
|
|
root_logger = logging.getLogger()
|
|
root_logger.setLevel(logging.DEBUG)
|
|
|
|
ch = logging.StreamHandler()
|
|
ch.setLevel(logging.DEBUG)
|
|
root_logger.addHandler(ch)
|
|
|
|
unittest.main(verbosity=2)
|
|
|