from requests_ratelimiter import LimiterSession
import os
import json
from tqdm import tqdm
import traceback
from time import sleep

def download_kernel_config(session, kernel_config_link, crash, bug_id) :
    kernel_config = None
    error_dict = {}
    try :
        kernel_config = session.get(kernel_config_link).text
    except Exception as e :
        error_dict = {bug_id: "kernel config could not be downloaded"}
        print(error_dict)
    return error_dict, kernel_config

def download_syz_reproducer(session, syz_repro_link, crash, bug_id) :
    syz_repro = None
    error_dict = {}
    try :
        syz_repro = session.get(syz_repro_link).text
        syz_repro_trimmed = syz_repro[syz_repro.find("#{"):]
    except Exception as e :
        error_dict = {bug_id: "syz-reproducer-data could not be downloaded"}
        print(error_dict)
    return error_dict, syz_repro_trimmed

def download_c_reproducer(session, c_repro_link, crash, bug_id) :
    c_repro = None
    error_dict = {}
    try :
        c_repro = session.get(c_repro_link).text
        c_repro_trimmed = c_repro[c_repro.find("// autogenerated by syzkaller"):]
    except Exception as e :
        error_dict = {bug_id: "c-reproducer-data could not be downloaded"}
        print(error_dict)
    return error_dict, c_repro_trimmed

def main(benchmark_path) :
    session = LimiterSession(per_second=1)
    bug_list = os.listdir(benchmark_path)

    for bug in tqdm(bug_list) :
        bug_path = os.path.join(benchmark_path,bug)
        bug_data = json.load(open(bug_path,"r"))
        bug_crashes = bug_data["crashes"]

        for crash in bug_crashes :
            try : 
                constant_prefx = "https://syzkaller.appspot.com"
                
                if crash.get("syz-reproducer") :
                    syz_repro_link =  constant_prefx + crash["syz-reproducer"]
                    _, syz_repro = download_syz_reproducer(session, syz_repro_link, crash, bug_data["id"])
                    sleep(1)

                if crash.get("c-reproducer") : 
                    c_repro_link = constant_prefx + crash["c-reproducer"]
                    _, c_repro = download_c_reproducer(session, c_repro_link, crash, bug_data["id"])
                    sleep(1)

                assert((crash.get("syz-reproducer") is not None) or (crash.get("c-reproducer") is not None))

                kernel_config_link = constant_prefx + crash["kernel-config"]
                _, kernel_config = download_kernel_config(session, kernel_config_link, crash, bug_data["id"])
                sleep(1.5)

                crash["syz-reproducer-data"] = syz_repro
                crash["c-reproducer-data"] = c_repro
                crash["kernel-config-data"] = kernel_config
            
            except Exception as e :
                print("Bug : {}".format(bug))
                print(traceback.print_exc())

            json.dump(bug_data,open(bug_path,"w"),indent=4)

if __name__ == '__main__' :
    main(benchmark_path= os.path.join(os.getenv("KBENCH_PATH")))