import subprocess
import lief
import os
import time
import re

PATH = "./coreutils-9.1/src/"

def init():
    out = open("out.csv", "w")
    out.write("Name,File size(KiB),Number of symbols,Number of imports,Restoration time(s),Execution time(s),File size(KiB),Number of symbols,Number of imports,Restoration time(s),Execution time(s)\n")
    return out

def replace_cmd(cmd, name):
    res = []
    for i in cmd:
        if i == "#":
            res.append(name)
        else:
            res.append(i)
    return res

def run_benchmark(file, name, cmd):
    print(f"[+] Running benchmark for {name} with command \"{cmd}\"")
    cmd = cmd.split(" ")
    norm_path = PATH + name
    obf_path = f"{PATH}{name}-dir/out/{name}-fixed"
    norm_size = int(os.path.getsize(norm_path) / 1024)
    obf_size = int(os.path.getsize(obf_path) / 1024)
    norm = lief.parse(norm_path)
    obf = lief.parse(obf_path)
    norm_symbols = len(norm.symbols)
    obf_symbols = 0
    for i in obf.symbols:
        if i.type != 0:
            obf_symbols += 1
    norm_imports = len(norm.imported_functions)
    obf_imports = 0
    for i in obf.imported_functions:
        if i.name != "":
            obf_imports += 1
    
    start = time.time()
    proc = subprocess.run(replace_cmd(cmd, norm_path), capture_output=True)
    end = time.time()
    norm_exe = end - start
    expect = proc.stdout
    norm_code = proc.returncode

    start = time.time()
    proc = subprocess.run(replace_cmd(cmd, obf_path), capture_output=True)
    end = time.time()
    obf_exe = end - start
    obf_out = proc.stdout
    obf_code = proc.returncode
    if obf_code == -11:
        print(f"\033[91m[!] Error in {name} (segfault)\033[0m")
        return
    if obf_code != norm_code: 
        print(f"\033[91m[!] Error in {name} (diff exit code)\033[0m")
        return
    match = re.search(b"restoration library time: ([0-9.]+)", obf_out)
    restore = float(match.group(1))
    
    out = f"{name},{norm_size},{norm_symbols},{norm_imports},N/A,{norm_exe:.3f},{obf_size},{obf_symbols},{obf_imports},{restore:.3f},{obf_exe:.3f}\n"
    if expect in obf_out:
        print("Results:", out)
        file.write(out)
    else:
        # Error in red
        print(f"\033[91m[!] Error in {name}\033[0m")
        # print(f"Expected: {expect}")
        # print(f"Got: {obf_out}")
        print("Results:", out)
        file.write(f"{name},Error\n")

# core="tee md5sum split cat shuf mkfifo pathchk runcon expand tty basename nice truncate echo du ptx join df pwd test csplit sort whoami touch dcgen unlink b2sum sleep fmt stty logname chgrp printenv seq uname sha224sum od date base64 realpath readlink dircolors timeout tac numfmt wc basenc comm nproc expr stdbuf cksum printf groups chcon factor tail env pr head kill uniq stat link make-prime-list sum tsort extract-magic mknod users dd who sha1sum mktemp cut sha256sum dir mkdir nl ginstall shred fold rmdir sha384sum mv dirname id base32 pinky ln hostid chroot ls true cp sync yes unexpand chown getlimits chmod uptime rm vdir false sha512sum tr paste nohup"

test_data = [
    ("md5sum", "# ./test_file.txt"), # fail
    ("split", "# ./test_file.txt /tmp/a"),
    ("cat", "# ./test_file.txt"),
    ("shuf", "# --random-source=./test_file.txt ./test_file.txt"),
    ("mkfifo", "# /tmp/a"),
    ("pathchk", "# ./test_file.txt"),
    ("expand", "# ./test_file.txt"),
    ("tty", "#"),
    ("basename", "# $PWD"),
    ("nice", "#"),
    ("truncate", "# -s 0 ./empty_file.txt"),
    ("echo", "# hello"),
    ("du", "# -h /tmp"),
    ("ptx", "# ./test_file.txt"),
    ("join", "# ./test_file.txt ./test_file.txt"),
    ("df", "# /etc"),
    ("pwd", "#"),
    ("test", "# -f ./test_file.txt"),
    ("csplit", "# ./test_file.txt 1"),
    ("sort", "# ./test_file.txt"),
    ("whoami", "#"),
    ("touch", "# /tmp/dd"),
    # ("dcgen", "#"), build fail
    ("unlink", "# /tmp/a"),
    ("b2sum", "# ./test_file.txt"),
    ("sleep", "# 1"),
    ("fmt", "# ./test_file.txt"),
    ("stty", "#"),
    ("logname", "#"),
    ("chgrp", "# root ./test_file.txt"),
    ("printenv", "#"),
    ("seq", "# 1 10"),
    ("uname", "#"),
    ("sha224sum", "# ./test_file.txt"), # fail segfault
    ("od", "# ./test_file.txt"),
    ("date", "#"),
    ("base64", "# ./test_file.txt"),
    ("realpath", "# ./test_file.txt"),
    ("readlink", "# ./test_file.txt"),
    ("dircolors", "#"),
    ("timeout", "# 1s sleep 2"),
    ("tac", "# ./test_file.txt"),
    ("numfmt", "# 1000"),
    ("wc", "# ./test_file.txt"),
    ("basenc", "# ./test_file.txt"),
    ("comm", "# ./test_file.txt ./test_file.txt"),
    ("nproc", "#"),
    ("expr", "# 1"),
    ("stdbuf", "#"), # need test case
    ("cksum", "# ./test_file.txt"),
    ("printf", "# hello"),
    ("groups", "#"),
    ("chcon", "# -t s0 ./test_file.txt"),
    ("factor", "# 10"),
    ("tail", "# -n 1 ./test_file.txt"),
    ("env", "#"),
    ("pr", "# ./test_file.txt"),
    ("head", "# -n 1 ./test_file.txt"),
    ("kill", "# $$"),
    ("uniq", "# ./test_file.txt"),
    ("stat", "# ./test_file.txt"),
    ("link", "# ./test_file.txt"), 
    ("make-prime-list", "# 10"), # fail
    ("sum", "# ./test_file.txt"),
    ("tsort", "# ./test_file.txt"),
    # ("extract-magic", "# ./test_file.txt"), build fail
    ("ls", "# -l"),
]
file = init()
for name, cmd in test_data:
    run_benchmark(file, name, cmd)

file.close()