import os
import useful
#
def check_path_traversal(path, base_path):
	# this is not guarenteed to be secure
	if os.path.commonprefix([os.path.realpath(path), base_path]) != base_path:
		raise Exception('The path \'{}\' is not in the base path \'{}.\''.format(os.path.realpath(path), base_path))
	#
#
def get_online_numa_nodes():
	with open('/sys/devices/system/node/online', 'r') as f:
		online_nodes = f.read().strip()
		return useful.parse_range_list(online_nodes)
	#
#
def get_processors_in_numa_node(node_id):
	path = '/sys/devices/system/node/node{}/cpulist'.format(node_id)
	check_path_traversal(path, '/sys/devices/system/node')
	with open(path, 'r') as f:
		return useful.parse_range_list(f.read().strip())
	#
#
def get_thread_siblings(processor_id):
	path = '/sys/devices/system/cpu/cpu{}/topology/thread_siblings_list'.format(processor_id)
	check_path_traversal(path, '/sys/devices/system/cpu')
	with open(path, 'r') as f:
		return useful.parse_range_list(f.read().strip())
	#
#
def get_numa_overview():
	numa_nodes = {}
	#
	for node_id in get_online_numa_nodes():
		numa_nodes[node_id] = {}
	#
	for node_id in numa_nodes:
		processors = get_processors_in_numa_node(node_id)
		#
		numa_nodes[node_id]['physical_cores'] = []
		for processor_id in processors:
			thread_siblings = sorted(get_thread_siblings(processor_id))
			#
			if thread_siblings not in numa_nodes[node_id]['physical_cores']:
				numa_nodes[node_id]['physical_cores'].append(thread_siblings)
			#
		#
	#
	return numa_nodes
#
if __name__ == '__main__':
	print(get_numa_overview())
#