mirror of
https://github.com/google/security-research
synced 2026-06-08 14:27:23 +00:00
Eduardo' Vela" Nava (sirdarckcat)
107 lines
3.1 KiB
Python
107 lines
3.1 KiB
Python
import sqlite3
|
|
from collections import defaultdict
|
|
|
|
|
|
def find_paths_to_function_recursive(db_file, function, max_depth=100):
|
|
"""
|
|
Finds paths to $function recursively, finding edges
|
|
by target message and filtering by rule_id.
|
|
"""
|
|
|
|
conn = sqlite3.connect(db_file)
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("SELECT id, message, uri, startLine FROM locations")
|
|
locations = cursor.fetchall()
|
|
|
|
# Create a dictionary to map location ID to (message, uri, startLine)
|
|
location_map = defaultdict(set)
|
|
for id, message, uri, startLine in locations:
|
|
location_map[(message, uri, startLine)].add(id)
|
|
|
|
# Find IDs corresponding to $function
|
|
function_keys = [
|
|
key for key, value in location_map.items() if key[0] == function
|
|
][0:1000]
|
|
|
|
def preprocess_edges(cursor):
|
|
cursor.execute(
|
|
"""
|
|
SELECT e.source_location_id, e.target_location_id
|
|
FROM edges AS e
|
|
WHERE e.rule_id = 'callgraph-all'
|
|
"""
|
|
)
|
|
edges = {}
|
|
for source_id, target_id in cursor.fetchall():
|
|
if target_id not in edges:
|
|
edges[target_id] = []
|
|
edges[target_id].append(source_id)
|
|
return edges
|
|
|
|
edges = preprocess_edges(cursor)
|
|
|
|
def find_edge_source(target_key):
|
|
target_ids = location_map.get(target_key)
|
|
if target_ids:
|
|
for target_id in target_ids:
|
|
source_ids = edges.get(target_id)
|
|
if source_ids:
|
|
return source_ids
|
|
return None
|
|
|
|
paths = [] # To store all unique paths
|
|
seen_paths = set() # To keep track of visited paths
|
|
|
|
def dfs_recursive(node_key, path, visited, depth):
|
|
visited.add(node_key)
|
|
path.append(node_key)
|
|
|
|
neighbors = find_edge_source(node_key)
|
|
|
|
if neighbors is not None:
|
|
for neighbor_id in neighbors:
|
|
for key, ids in location_map.items():
|
|
if neighbor_id in ids:
|
|
neighbor_key = key
|
|
break
|
|
if (
|
|
neighbor_key
|
|
and depth < max_depth
|
|
and neighbor_key not in visited
|
|
):
|
|
dfs_recursive(
|
|
neighbor_key, path.copy(), visited.copy(), depth + 1
|
|
)
|
|
else:
|
|
path_tuple = tuple(path)
|
|
if path_tuple not in seen_paths:
|
|
seen_paths.add(path_tuple)
|
|
paths.append(path.copy())
|
|
|
|
path.pop()
|
|
visited.remove(node_key)
|
|
|
|
for start_key in function_keys:
|
|
dfs_recursive(start_key, [], set(), 0)
|
|
|
|
named_paths = []
|
|
for path in paths:
|
|
named_paths.append(path.copy())
|
|
|
|
conn.close()
|
|
return named_paths
|
|
|
|
|
|
if __name__ == "__main__":
|
|
db_file = "codeql_data-6.1.db"
|
|
paths = find_paths_to_function_recursive(db_file, "___sys_recvmsg")
|
|
|
|
for path in paths:
|
|
# Format the path
|
|
formatted_path = " -> ".join(
|
|
[f"{node[0]} ({node[1]}:{node[2]})" for node in path]
|
|
)
|
|
print(formatted_path)
|
|
print()
|