# import unittest
import subprocess, resource
import lief
import os
# import time
import re

PATH = "./coreutils-9.1/src"

class Line:
    file = None
    @classmethod
    def init(cls):
        cls.file = open("out.csv", "w")
        cls.file.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")
    
    def write(self):
        out = f"{self.name},{self.norm_size},{self.norm_symbols},{self.norm_imports},0,{self.norm_exe:.3f},{self.obf_size},{self.obf_symbols},{self.obf_imports},{self.restore:.3f},{self.obf_exe:.3f}\n"
        self.file.write(out)

    def __init__(self, name: str) -> None:
        self.name = name
        self.norm_path = f"{PATH}/{name}"
        self.obf_path = f"{PATH}/{name}-dir/out/{name}-fixed"

def numOfSymbols(sym: list) -> int:
    count = 0
    for i in sym:
        if i.type != 0:
            count += 1
    return count

def numOfImports(imp: list) -> int:
    count = 0
    for i in imp:
        if i.name != "":
            count += 1
    return count

def setup(binpath: str, libpath: str = None):
    if os.path.isdir("/tmp/test"):
        os.system("rm -rf /tmp/test")
    os.mkdir("/tmp/test")
    # if libpath:
    #     os.system(f"cp {binpath} {libpath} ./test_file.txt /tmp/test")
    # else: 
    #     os.system(f"cp {binpath} ./test_file.txt /tmp/test")
    os.system("cp ./test_file.txt /tmp/test")


# class Benchmark(unittest.TestCase):
def info(l: Line):
    l.norm_size = int(os.path.getsize(l.norm_path) / 1024)
    l.obf_size = int(os.path.getsize(l.obf_path) / 1024)

    norm = lief.parse(l.norm_path)
    obf = lief.parse(l.obf_path)

    l.norm_symbols = numOfSymbols(norm.symbols)
    l.obf_symbols = numOfSymbols(obf.symbols)
    l.norm_imports = numOfImports(norm.imported_functions)
    l.obf_imports = numOfImports(obf.imported_functions)

def run(l, cmd):
    print(f"[+] Running benchmark for {l.name} with command \"{cmd}\"")
    cmd = cmd.split(" ")
    setup(l.norm_path)
    start = resource.getrusage(resource.RUSAGE_CHILDREN).ru_utime
    p1 = subprocess.run([l.norm_path] + cmd, capture_output=True)
    end = resource.getrusage(resource.RUSAGE_CHILDREN).ru_utime
    l.norm_exe = end - start

    setup(l.obf_path)
    start = resource.getrusage(resource.RUSAGE_CHILDREN).ru_utime
    p2 = subprocess.run([l.obf_path] + cmd, capture_output=True)
    end = resource.getrusage(resource.RUSAGE_CHILDREN).ru_utime
    if p2.returncode == -11:
        print(f"\033[91m[!] Error in {l.name} (segfault)\033[0m")
        Line.file.write(f"{l.name},segfault\n")
        return
    # print(p2.stdout)
    match = re.search(b"restoration library time: ([0-9.]+)", p2.stdout)
    l.restore = float(match.group(1))
    l.obf_exe = end - start
    if p2.returncode != p1.returncode:
        print(f"\033[91m[!] Error in {l.name} (diff exit code)\033[0m")
        Line.file.write(f"{l.name},exit code diff\n")
        return
    l.write()
    # if p1.stdout in p2.stdout:
    #     l.write()
    # else:
    #     print(f"\033[91m[!] Error in {l.name} (stdout diff)\033[0m")
    #     print(p1.stdout)
    #     print("-"*20)
    #     print(p2.stdout)

    #     Line.file.write(f"{l.name},stdout diff\n")
    

def test_basic(name, cmd):
    l = Line(name)
    info(l)
    run(l, cmd)

test_data = [
    ("md5sum", "/tmp/test/test_file.txt"),
    ("split", "/tmp/test/test_file.txt /tmp/test/out"),
    ("cat", "/tmp/test/test_file.txt"),
    ("mkfifo", "/tmp/test/a"),
    ("shuf", "--random-source=/tmp/test/test_file.txt /tmp/test/test_file.txt"),
    ("pathchk", "/tmp/test/test_file.txt"),
    ("expand", "/tmp/test/test_file.txt"),
    ("tty", ""),
    ("basename", "/tmp/test/test_file.txt"),
    ("nice", ""),
    ("truncate", "-s 0 /tmp/test/test_file.txt"),
    ("echo", "hello"),
    ("du", "-h /tmp"),
    ("ptx", "/tmp/test/test_file.txt"),
    ("join", "/tmp/test/test_file.txt /tmp/test/test_file.txt"),
    ("df", "--help"),
    ("pwd", ""),
    ("test", "-f /tmp/test_file.txt"),
    ("csplit", "/tmp/test_file.txt 1"),
    ("sort", "/tmp/test_file.txt"),
    ("whoami", ""),
    ("touch", "/tmp/test/a"),
    ("unlink", "/tmp/test/test_file.txt"),
    ("b2sum", "/tmp/test/test_file.txt"),
    ("sleep", "1"),
    ("fmt", "/tmp/test/test_file.txt"),
    ("stty", ""),
    ("logname", ""),
    ("chgrp", "root /tmp/test/test_file.txt"),
    ("printenv", ""),
    ("seq", "1 10"),
    ("uname", ""),
    ("sha224sum", "/tmp/test/test_file.txt"),
    ("od", "/tmp/test/test_file.txt"),
    ("date", ""),
    ("base64", "/tmp/test/test_file.txt"),
    ("realpath", "/tmp/test/test_file.txt"),
    ("readlink", "/tmp/test/test_file.txt"),
    ("dircolors", ""),
    ("timeout", "1s sleep 2"),
    ("tac", "/tmp/test/test_file.txt"),
    ("numfmt", "1000"),
    ("wc", "/tmp/test/test_file.txt"),
    ("basenc", "/tmp/test/test_file.txt"),
    ("comm", "/tmp/test/test_file.txt /tmp/test/test_file.txt"),
    ("nproc", ""),
    ("expr", "1"),
    ("cksum", "/tmp/test/test_file.txt"),
    ("printf", "hello"),
    ("groups", ""),
    ("chcon", "-t s0 /tmp/test/test_file.txt"),
    ("factor", "10"),
    ("tail", "-n 1 /tmp/test/test_file.txt"),
    ("env", ""),
    ("pr", "/tmp/test/test_file.txt"),
    ("head", "-n 1 /tmp/test/test_file.txt"),
    ("kill", "$$"),
    ("uniq", "/tmp/test/test_file.txt"),
    ("stat", "-f /tmp/test/test_file.txt"),
    ("link", "/tmp/test/test_file.txt /tmp/test/test_file.txt"),
    # ("make-prime-list", "10"), # build fail
    ("sum", "/tmp/test/test_file.txt"),
    ("tsort", "/tmp/test/test_file.txt"),
    # ("extract-magic", "/tmp/test/test_file.txt"), build fail
    ("mknod", "/tmp/test/test_file.txt"),
    ("users", ""),
    ("dd", "--help"),
    ("who", ""),
    ("sha1sum", "/tmp/test/test_file.txt"),
    ("mktemp", ""),
    ("cut", "-c 1 /tmp/test/test_file.txt"),
    ("sha256sum", "/tmp/test/test_file.txt"),
    ("dir", "/tmp/test/test_file.txt"),
    ("mkdir", "/tmp/test/a"),
    ("nl", "/tmp/test/test_file.txt"),
    ("ginstall", "/tmp/test/test_file.txt /tmp/test/test_file.txt"),
    ("shred", "-u /tmp/test/test_file.txt"),
    ("fold", "-w 10 /tmp/test/test_file.txt"),
    ("rmdir", "/tmp/test/a"),
    ("sha384sum", "/tmp/test/test_file.txt"),
    ("mv", "/tmp/test/test_file.txt /tmp/test/test_file.txt"),
    ("dirname", "/tmp/test/test_file.txt"),
    ("id", ""),
    ("base32", "/tmp/test/test_file.txt"),
    ("pinky", ""),
    ("ln", "/tmp/test/test_file.txt /tmp/test/test_file.txt"),
    ("hostid", ""),
    ("chroot", "/tmp/test /tmp/test/test_file.txt"),
    ("ls", "/tmp/test"),
    ("true", ""),
    ("cp", "/tmp/test/test_file.txt /tmp/test/test_file.txt"),
    ("sync", ""),
    ("yes", "--help"),
    ("unexpand", "/tmp/test/test_file.txt"),
    ("chown", "root /tmp/test/test_file.txt"),
    ("getlimits", ""),
    ("chmod", "777 /tmp/test/test_file.txt"),
    ("uptime", ""),
    ("rm", "/tmp/test/test_file.txt"),
    ("vdir", "/tmp/test"),
    ("false", ""),
    ("sha512sum", "/tmp/test/test_file.txt"),
    ("tr", "a b /tmp/test/test_file.txt"),
    ("paste", "/tmp/test/test_file.txt /tmp/test/test_file.txt"),
    ("nohup", "sleep 1")
]

# 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"

if __name__ == "__main__":
    Line.init()
    for name, cmd in test_data:
        test_basic(name, cmd)
    # unittest.main()