search_objects.py

import pandas as pd
import json
import os
import re
from config import STATIC_DATA_DIR, trace_dir

ast_path = ""
lock_struct_mapping_db = ""


def traverse_ast(ast_json):
    if isinstance(ast_json, dict):
        if 'detail' in ast_json and ast_json['detail'] == 'lock':
            # pretty print json 
            print(json.dumps(ast_json, indent=4))
            print(ast_json['range'])
            # print(ast_json)
            # If we encounter a node with detail "lock", we stop traversing
            return True
        elif 'children' in ast_json:
            # Recursively traverse through the children
            for child in ast_json['children']:
                if traverse_ast(child):
                    return True  # Propagate up if a lock is found in children
    elif isinstance(ast_json, list):
        # Recursively traverse through each element of the list
        for item in ast_json:
            if traverse_ast(item):
                return True  # Propagate up if a lock is found in children
    return False


def map_objects_test():
    # load the json file
    working_dir = os.getcwd()
    ast_out_path = os.path.join(working_dir, 'ast')
    ast_file = "fs_proc_base.c.json"
    with open(ast_out_path + '/' + ast_file) as f:
        data = json.load(f)

    # filter_func = "check_unsafe_exec"
    filter_func = "proc_pid_make_base_inode"
    count = 0 

    # print(type(data))

    #print length of data
    # print(len(data))
    # iterate over five items in data
    # print(data[0])

    # top key has 4 things: arcana, children, kind, role
    for key in data:
        # print(key)
        # print(type(item))
        # print(item)
        count += 1
        acquire_line = 163
        # children contains the contents of the file
        if key == 'children':
            # print(key)
            # print(len(data[key][0]))
            for item in data[key]:

                # check if detail and kind are in the child
                if "detail" in item and "kind" in item and "range" in item:
                    # print(child["detail"])
                    # print(child["kind"])
                    if item["detail"] == filter_func:
                        print(len(item["children"]))
                        # if(len(item["children"]) != 2):
                        #     print(item["detail"])
                

                        # iterate over the components of the function
                        for child in item["children"]:
                            # print(child)
                            # print()

                            # check if arcana in item starts with CompoundStmt, this is the function body
                            # print(child["arcana"])
                            if "CompoundStmt" in child["arcana"]:
                                # call the traverse_ast function to check if there is a lock in the function body
                                if traverse_ast(child):
                                    print("Lock found in function body")
                                    # print(item["detail"])

                                # iterate over the lines in the function body
                                # for line in child["children"]:
                                #     # print(line["arcana"])
                                #     # ignore DeclStmt for now 
                                #     if "DeclStmt" in line["arcana"]:
                                #         continue
                                    

                                #     # checke for ForStmt
                                #     if "ForStmt" in line["arcana"]:
                                #         print(len(line["children"]))
                                #         for stmt in line["children"]:
                                #             print(stmt["arcana"])

                                        # recursively parse each child of the ForStmt, we don't know how many children it has                                        
                                            
                                    # check if the line has loops or condirional stattements like ForStmt, IfStmt, WhileStmt



                        

                        # separate signature and body (compound statement)

                # print(child["detail"])
                # print(child["kind"])
                # print(type(child))
                
                # count += 1
                # if count == 10:
                #     break
            
        else:
            # print(data[item])
            continue
        # if count == 1:
        #     break

def map_objects(search_file):

    # load unique lock struct map
    lock_struct_mapping_db = STATIC_DATA_DIR / "lock_struct_map_db" / "lock_struct_map.csv"
    lock_struct_map = pd.read_csv(lock_struct_mapping_db)

    # load drivers db
    drivers_lock_struct_map_db = pd.read_csv(STATIC_DATA_DIR / "lock_struct_map_db" / "lock_struct_map_drivers.csv")

    # concatenate lock_struct_map and drivers_lock_struct_map_db
    # lock_struct_map = pd.concat([lock_struct_map, drivers_lock_struct_map_db])
    
    # covert file, type, name, lock_type,lock_name columns to string
    lock_struct_map["file"] = lock_struct_map["file"].astype(str)
    lock_struct_map["type"] = lock_struct_map["type"].astype(str)
    lock_struct_map["name"] = lock_struct_map["name"].astype(str)
    lock_struct_map["lock_type"] = lock_struct_map["lock_type"].astype(str)
    lock_struct_map["lock_name"] = lock_struct_map["lock_name"].astype(str)

    # do the same column conversion for drivers_lock_struct_map_db
    drivers_lock_struct_map_db["file"] = drivers_lock_struct_map_db["file"].astype(str)
    drivers_lock_struct_map_db["type"] = drivers_lock_struct_map_db["type"].astype(str)
    drivers_lock_struct_map_db["name"] = drivers_lock_struct_map_db["name"].astype(str)
    drivers_lock_struct_map_db["lock_type"] = drivers_lock_struct_map_db["lock_type"].astype(str)
    drivers_lock_struct_map_db["lock_name"] = drivers_lock_struct_map_db["lock_name"].astype(str)


    # load global lock names
    global_locks = pd.read_csv(STATIC_DATA_DIR / "global_locks" / "global_locks.csv")
    global_locks["lock_name"] = global_locks["lock_name"].astype(str)

    # print type of each column in lock_struct_map
    # for col in lock_struct_map.columns:
    #     print(col, lock_struct_map[col].dtype)

    # load the search file
    lock_details = pd.read_csv(search_file)

    # covert lock_name to string
    lock_details["LockName"] = lock_details["LockName"].astype(str)

    # change columns EndFunction, EndFile to string
    lock_details["EndFunction"] = lock_details["EndFunction"].astype(str)
    lock_details["EndFile"] = lock_details["EndFile"].astype(str)
    
    # add a column to store the lock struct and struct path 
    lock_details["Global"] = 0

    # lock_details["StructPath"] = ""
    # count = 0
    # iterate over the rows in the search file
    for index, row in lock_details.iterrows():
        # print processing ow details
        print("Processing row: ", index)
        print("Lock Name: ", row["LockName"])
        # count += 1

        # if count == 4:
        #     break

        # if row["FromLineNumber"] == -1:
        #     continue

        lock_name = row["LockName"]
        lock_primitive = row["LockPrimitive"]
        file = row["EndFile"]

        # if "[" in lock_name:
        #     print("Lock Name: ", lock_name)

        # check if lock_name has [], if so remove along with the content in between [ and  ]
        lock_name = re.sub(r'\[.*?\]', '', lock_name)
        
        print("Lock Name: ", lock_name)

        
        # Step 1: Match with global locks
        # check if the lock is a global lock
        if "&" in lock_name: # remove & and then match with global locks
             if lock_name.split("&")[1] in global_locks["lock_name"].values:
                lock_details.at[index, "Global"] = 1
                print("Match found in global locks")
                continue
        

        if lock_name in global_locks["lock_name"].values:
            lock_details.at[index, "Global"] = 1
            print("Match found in global locks")
            continue

        # if there are more than one "->"" in the lock name eg: &(&sig->stats_lock)->lock, skip 
        if lock_name.count("->") > 1:
            continue
        
        if "->" in lock_name:
            lock_name = lock_name.split("->")[1]
        
        if "." in lock_name: #eg: &htab->buckets[i].raw_lock
            lock_name = lock_name.split(".")[1]

        lock_name = lock_name.strip()
        

        # if ")" in lock_name:
        #     lock_name = lock_name.split(")")[0]
        
        # print("Lock Name: ", lock_name)
 
        # Step 2: Match with lock_struct_map
        # check for lock_name substring in lock_struct_map
        # lock_struct = lock_struct_map[lock_struct_map["lock_name"].str.contains(lock_name)]
        lock_struct = lock_struct_map[lock_struct_map["lock_name"].str.contains(re.escape(lock_name))]


        # if lock_name == "file_lock":
        #     print(lock_struct)
            # print(len(lock_struct))
        
        # find exact match for lock name 
        if len(lock_struct) > 1:
            # print("Multiple matches found for lock: ", lock_name)
            # print(lock_struct)
            # find the exact match
            lock_struct = lock_struct[lock_struct["lock_name"] == lock_name]
            # print(lock_struct)

            #TODO: find match for lock type
            if len(lock_struct) > 1:
                if "spin" in lock_primitive:
                    lock_struct = lock_struct[(lock_struct["lock_name"] == lock_name) & (lock_struct["lock_type"] == "spinlock_t")]
                elif "mutex" in lock_primitive:
                    lock_struct = lock_struct[(lock_struct["lock_name"] == lock_name) & (lock_struct["lock_type"] == "struct mutex")]
                elif "rw" in lock_primitive:
                    lock_struct = lock_struct[(lock_struct["lock_name"] == lock_name) & (lock_struct["lock_type"] == "rwlock_t")]
                elif "sem" in lock_primitive:
                    lock_struct = lock_struct[(lock_struct["lock_name"] == lock_name) & (lock_struct["lock_type"] == "struct rw_semaphore")]


        # if match is found and is unique
        if len(lock_struct) == 1:
            print("Match found for lock: ", lock_name)
            lock_struct = lock_struct.iloc[0]
            lock_details.at[index, "LockStruct"] = lock_struct["name"]
            lock_details.at[index, "StructPath"] = lock_struct["file"]
        
        # check for driver when match is not found in the lock_struct_map
        else:
            # print("No match found for lock: ", lock_name)
            if file.startswith("drivers/"):
                # do the same checks with drivers_lock_struct_map_db
                lock_struct = drivers_lock_struct_map_db[drivers_lock_struct_map_db["lock_name"].str.contains(re.escape(lock_name))]
                if len(lock_struct) > 1:
                    # print("Multiple matches found for lock for driver check: ", lock_name)
                    # print(lock_struct[(lock_struct["lock_name"] == lock_name) & (lock_struct["type"].astype(str) == "struct mutex")])
                    # print(lock_struct)
                   
                    # find the exact match
                    # lock_struct = lock_struct[lock_struct["lock_name"] == lock_name]
                    # print(lock_struct)

                    # match for lock type
                    # print("lock_primitive: ", lock_primitive)   
                    if "spin" in lock_primitive:
                        lock_struct = lock_struct[(lock_struct["lock_name"] == lock_name) & (lock_struct["lock_type"] == "spinlock_t")]
                    elif "mutex" in lock_primitive:
                        # print("Mutex in lock primitive")
                        lock_struct = lock_struct[(lock_struct["lock_name"] == lock_name) & (lock_struct["lock_type"] == "struct mutex")]
                    elif "rw" in lock_primitive:
                        lock_struct = lock_struct[(lock_struct["lock_name"] == lock_name) & (lock_struct["lock_type"] == "rwlock_t")]
                    elif "sem" in lock_primitive:
                        lock_struct = lock_struct[(lock_struct["lock_name"] == lock_name) & (lock_struct["lock_type"] == "struct rw_semaphore")]
                
                # print("Lock Struct: ", lock_struct)
                if len(lock_struct) == 1:
                    print("Match found in driver for lock: ", lock_name)
                    lock_struct = lock_struct.iloc[0]
                    lock_details.at[index, "LockStruct"] = lock_struct["name"]
                    lock_details.at[index, "StructPath"] = lock_struct["file"]
            
            continue
    
    # save the updated file
    lock_details.to_csv(search_file, index=False)
    

def map_lock_defs_objects(search_file):

    # load generic and non-generic lock struct map
    generic_lock_struct_mapping_db = pd.read_csv(STATIC_DATA_DIR / "lock_defs" / "generic_locks_defs.csv")
    non_generic_lock_struct_mapping_db = pd.read_csv(STATIC_DATA_DIR / "lock_defs" / "non_generic_locks_defs.csv")
 
    
    generic_lock_struct_mapping_db["from_line"] = generic_lock_struct_mapping_db["from_line"].astype(int)

    # convert func,path,acquire_function,lock_name,lock_def_path,struct_name to string
    generic_lock_struct_mapping_db["func"] = generic_lock_struct_mapping_db["func"].astype(str).str.strip().str.lower()
    generic_lock_struct_mapping_db["path"] = generic_lock_struct_mapping_db["path"].astype(str).str.strip().str.lower()
    generic_lock_struct_mapping_db["acquire_function"] = generic_lock_struct_mapping_db["acquire_function"].astype(str).str.strip().str.lower()
    generic_lock_struct_mapping_db["lock_name"] = generic_lock_struct_mapping_db["lock_name"].astype(str).str.strip().str.lower()
    generic_lock_struct_mapping_db["lock_def_path"] = generic_lock_struct_mapping_db["lock_def_path"].astype(str).str.strip().str.lower()
    generic_lock_struct_mapping_db["struct_name"] = generic_lock_struct_mapping_db["struct_name"].astype(str).str.strip().str.lower()

    # do the same column conversion for non_generic_lock_struct_mapping_db
    non_generic_lock_struct_mapping_db["from_line"] = non_generic_lock_struct_mapping_db["from_line"].astype(int)
    non_generic_lock_struct_mapping_db["func"] = non_generic_lock_struct_mapping_db["func"].astype(str).str.strip().str.lower()
    non_generic_lock_struct_mapping_db["path"] = non_generic_lock_struct_mapping_db["path"].astype(str).str.strip().str.lower()
    non_generic_lock_struct_mapping_db["acquire_function"] = non_generic_lock_struct_mapping_db["acquire_function"].astype(str).str.strip().str.lower()
    non_generic_lock_struct_mapping_db["lock_name"] = non_generic_lock_struct_mapping_db["lock_name"].astype(str).str.strip().str.lower()
    non_generic_lock_struct_mapping_db["lock_def_path"] = non_generic_lock_struct_mapping_db["lock_def_path"].astype(str).str.strip().str.lower()
    non_generic_lock_struct_mapping_db["struct_name"] = non_generic_lock_struct_mapping_db["struct_name"].astype(str).str.strip().str.lower()

    # load the search file
    lock_details_df = pd.read_csv(search_file)

    # change columns EndFunction, EndFile to string
    lock_details_df["FromLineNumber"] = lock_details_df["FromLineNumber"].astype(int)
    lock_details_df["EndFunction"] = lock_details_df["EndFunction"].astype(str).str.strip().str.lower()
    lock_details_df["EndFile"] = lock_details_df["EndFile"].astype(str).str.strip().str.lower()

    # # print types of each dataframe
    # print(lock_details_df.dtypes)
    # print(generic_lock_struct_mapping_db.dtypes)
    # print(non_generic_lock_struct_mapping_db.dtypes)

    for index, row in lock_details_df.iterrows():
        struct_name = row["LockStruct"]

        print("LockStruct: ", struct_name)

        # check if LockStruct is nan
        if pd.isnull(struct_name):
            # print("LockStruct is empty")
            func = row["EndFunction"].strip().lower()
            file = row["EndFile"].strip().lower()
            line = row["FromLineNumber"]

            print("Func: ", type(line))
            # print("File: ", file)

            # check if the function is in lock_struct_mapping_db where func and file match and line with either equal, +- 1 equal with in(from_line)
            # check 1: check in generic lock struct mapping db
            lock_struct = generic_lock_struct_mapping_db[(generic_lock_struct_mapping_db["func"] == func) & (generic_lock_struct_mapping_db["path"] == file)& ((generic_lock_struct_mapping_db["from_line"] == line) | (generic_lock_struct_mapping_db["from_line"] == line - 1) | (generic_lock_struct_mapping_db["from_line"] == line + 1))]
            # print(lock_struct)
            if len(lock_struct) == 1:
                lock_struct = lock_struct.iloc[0]
                lock_details_df.at[index, "LockStruct"] = lock_struct["struct_name"]
                lock_details_df.at[index, "StructPath"] = lock_struct["lock_def_path"]
                print("LockStruct 1: ", lock_struct["struct_name"])
            
            else:
                print("No match found in generic lock struct mapping db")
                # check 2: check in non_generic lock struct mapping db
                lock_struct = non_generic_lock_struct_mapping_db[(non_generic_lock_struct_mapping_db["func"] == func) & (non_generic_lock_struct_mapping_db["path"] == file) & ((non_generic_lock_struct_mapping_db["from_line"] == line) | (non_generic_lock_struct_mapping_db["from_line"] == line - 1) | (non_generic_lock_struct_mapping_db["from_line"] == line + 1))]
                if len(lock_struct) == 1:
                    lock_struct = lock_struct.iloc[0]
                    lock_details_df.at[index, "LockStruct"] = lock_struct["struct_name"]
                    lock_details_df.at[index, "StructPath"] = lock_struct["lock_def_path"]
                    print("LockStruct 2: ", lock_struct["struct_name"])


        else:
            continue

    # save the updated file
    lock_details_df.to_csv(search_file, index=False)
    return 

# define main function
if __name__ == "__main__":

    file = trace_dir + "ubuntu22/global_processed_traces/combined_lock_struct_map_v2.csv"
    # call seach objects function
    # map_objects(ast_path, lock_struct_mapping_db)
    map_objects(file)
    print("Completed mapping lock objects")
    map_lock_defs_objects(file)
    # map_objects("/mydata/locks-access-traces/global_processed_traces/test_driver.csv")
    # map_lock_defs_objects("/mydata/locks-access-traces/global_processed_traces/traces_lock_struct_map_v1.csv")