shell-test.py 2.46 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
from pwn import *
import time

# timeout is in seconds, 1s would be enough for the "default" case
def test_cmd(cmd, expstr=[], lines=1, timeout=1, mintime=None):
	# print(cmd)
	p.sendline(cmd)

	t0 = time.time()
	if expstr != []:
		outstr = p.readlines(numlines=lines, timeout=timeout)

	t1 = time.time()
	# discard the rest of the output
	p.clean(0)

	# verify whether the command finished too early
	if mintime and (t1 - t0) < mintime:
		log.failure(f"Test for {cmd} failed: The command finished earlier than expected.")
		return False


	# expect empty result
	if expstr == []:
		log.success(f"Test for {cmd} passed.")
		return True

	# output number of lines is different
	if len(expstr) != len(outstr):
		log.failure(f"Test for {cmd} failed: Expected {expstr} but got {outstr}")
		return False

	for ous, exs in zip(outstr, expstr):
		ous = ous.strip().decode()
		if (exs not in ous):
			log.failure(f"Test for {cmd} failed: Expected {exs} but got {ous}")
			return False
	log.success(f"Test for {cmd} passed.")
	return True

if __name__ == "__main__":
	if len(sys.argv) != 2:
		log.error(f"Run as f{sys.argv[0]} SHELL_EXEC")
	else:
		shell_exec = os.path.abspath(sys.argv[1])
		with tempfile.TemporaryDirectory() as tmpdir:
			os.chdir(tmpdir)
			log.info(f"Working in {tmpdir}")
			elf = ELF(shell_exec, checksec=False)
			p = elf.process()

			test_cmd('pwd', [tmpdir])
			test_cmd('/bin/pwd', [tmpdir])
			test_cmd('echo Hello!', ['Hello!'])
			test_cmd('ls -a',['.  ..'])
			test_cmd('touch file1.txt')
			test_cmd('touch file1.txt; ls', ['file1.txt'])
			test_cmd('echo world &',['world'])
			# test_cmd('echo hello && echo world',['hello','world'],lines=2)
			# test_cmd('sleep 3 && echo world &',['world'],lines=2,timeout=1)
			test_cmd('sleep 3 ; echo hello', ['hello'],timeout=4,mintime=3)
			test_cmd('echo Hi > file2.txt')
			test_cmd('cat < file2.txt', ['Hi'])
			test_cmd('cat < file2.txt > file3.txt; cat file3.txt',['Hi'])
			test_cmd('cat file2.txt | wc -l',['1'])
			test_cmd('cat file2.txt | cat | cat | cat | cat | wc -l',['1'])
			test_cmd('echo 1 && echo 2',['1','2'], lines=2)
			test_cmd('echo x | cat | wc -l && echo 2', ['1','2'], lines=2)
			test_cmd('sleep 2; echo Done',['Done'],timeout=3)
			test_cmd('cd /home; cd -; pwd',[tmpdir,tmpdir], lines=2)
			test_cmd('cd /home; cd -; cd -; cd -',[tmpdir,'/home',tmpdir], lines=3)
			test_cmd('cd .; pwd',[tmpdir])
			test_cmd(f'cd /home; cd ..; pwd; cd {tmpdir}',['/'])
			test_cmd(f'cd /home; cd ..; cd ..; cd ..; pwd; cd {tmpdir}',['/'])