import gzip
import inspect
import os
from io import StringIO
from typing import Optional, List

import numpy as np
import pandas as pd


def time_map(time_a: float, time_b: float, packet_a: int, packet_b: int, time_c: int, window_tolerance: int = 0) -> \
		Optional[float]:
	"""
	Map an API time into a packet number.
	This function was done in order to have a nice visualisation.
	The window tolereance is used to capture nearby calls
			=========     |
			ta      tb    tc=>pc
			pa      pb

	Args:
		time_a (float): time of the begining of the flow (in seconds since epoch, time.time())
		time_b (float): time of the ending of the flow
		packet_a (int): packet number of the begining of the flow
		packet_b (int): packet number of the ending of the flow
		time_c (float): time of the api call
		window_tolerance (int): time shift in second in which api calls are still considered to belong to the flow

	Returns:
		int: packet number, return None if the mapping fails
	"""
	
	window_tolerance *= 1000000
	# Check if inside flow
	if time_a <= time_c <= time_b:  # Chain comparaison
		return packet_a + (time_c - time_a) * (packet_b - packet_a) / (time_b - time_a)
	# Check if in the window border (simple)
	if time_a - window_tolerance <= time_c <= time_b:
		return packet_a
	"""if time_a <= time_c <= time_b + window_tolerance:
		return packet_b"""
	# Outside
	return None


def pure_time_map(time_a: float, time_b: float, time_c: float, window_tolerance: int = 0) -> Optional[float]:
	"""
	Map an API time into a packet number.
	This function was done in order to have a nice visualisation.
	The window tolereance is used to capture nearby calls
			=========     |
			ta      tb    tc=>pc

	Args:
		time_a (float): time of the begining of the flow (in seconds since epoch, time.time())
		time_b (float): time of the ending of the flow
		time_c (float): time of the api call
		window_tolerance (int): time shift in second in which api calls are still considered to belong to the flow

	Returns:
		int: time, return None if the mapping fails
	"""
	
	window_tolerance *= 1000000
	# Check if inside flow
	if time_a <= time_c <= time_b:  # Chain comparaison
		return time_c
	# Check if in the window border (simple)
	if time_a - window_tolerance <= time_c <= time_b:
		return time_a
	"""if time_a <= time_c <= time_b + window_tolerance:
		return time_b"""
	# Outside
	return None


def get_child_pids(current_pid: int, diff: pd.DataFrame) -> List:
	"""
	Get the child process pid of one process given its pid

	Args:
		current_pid (int): parent pid
		diff (pd dataframe): df recorded

	Returns:
		list: list of pids
	"""
	
	childs = diff[diff["parent_pid"] == current_pid]
	if childs.shape[0] == 0:
		return [current_pid]
	return [current_pid] + [v for index, row in childs.iterrows() for v in get_child_pids(row["process_id"], diff)]


def get_malware_pids(malware_name: str = "2020-09-30-Trickbot-EXE.exe", path: str = "./") -> List:
	"""
	Get the pids of all the malware generated processes

	Args:
		malware_name (str, optional): name of the malware. Defaults to "2020-09-08-Trickbot-EXE-gtag-ono72.exe".

	Returns:
		list: list of pids
		path: path of the malware
	"""
	
	first = pd.read_csv(path + "process_pre.csv")
	post = pd.read_csv(path + "process_post.csv")
	
	first.drop(first.columns[0], axis=1, inplace=True)
	post.drop(post.columns[0], axis=1, inplace=True)
	
	diff = first.merge(post, indicator=True,
	                   how='right').loc[lambda x: x['_merge'] != 'both']
	try:
		malware_pid = int(diff.loc[diff['process_name'] == malware_name]["process_id"].astype(float))
	except TypeError:
		raise RuntimeError('Malware PID not found, check malware name')
	return get_child_pids(malware_pid, diff)


def gzip_to_string(file_path: str) -> str:
	"""
	Open a gzip file and load the content in a string
	This function exists because the gzip may not be proprerly closed.
	In this case, the end is corrupted but the rest can be read.
	Args:
		file_path (string): path of the gzip file

	Returns:
		str: content of the gzip file
	"""
	
	gzip_file = gzip.open(file_path, "rt")
	string = ""
	while True:
		try:
			line = gzip_file.readline()
			string += line
		except EOFError:
			break
	return string


def get_malware_traces(path: str = "./") -> List:
	"""
	Get a list of dataframe representing the frida trace
	Its current format is [time, api_name, category]

	Returns:
		list: list of dataframe
	"""
	
	pids = list(set(get_malware_pids(path=path)))
	traces = []
	for pid in pids:
		if not os.path.isfile(f"{path}frida_{pid}.txt.gz"):
			print("Trace for {pid} does not exist")
		header = ["time", "api", "category"]
		frida_str = StringIO(gzip_to_string(f"{path}frida_{pid}.txt.gz"))
		dataframe = pd.read_csv(frida_str, names=header)
		dataframe.drop(
			dataframe.loc[dataframe['api'] == 'error'].index, inplace=True)
		dataframe.drop_duplicates(
			subset=['time', 'api'], keep='first', inplace=True)
		dataframe.reset_index(drop=True, inplace=True)
		traces.append(dataframe)
		dataframe["time_int"] = (dataframe["time"] * 1000000).astype(int)
	return traces


class Singleton(type):
	# Singleton modified to handle arguments (singleton for each argument set)
	_instances = {}
	_init = {}
	
	def __init__(cls, name, bases, dct):
		cls._init[cls] = dct.get('__init__', None)
	
	def __call__(cls, *args, **kwargs):
		init = cls._init[cls]
		if init is not None:
			key = (cls, frozenset(inspect.getcallargs(init, None, *args, **kwargs).items()))
		else:
			key = cls
		
		if key not in cls._instances:
			cls._instances[key] = super(Singleton, cls).__call__(*args, **kwargs)
		return cls._instances[key]


class MalwareTraceExtractor(metaclass=Singleton):
	def __init__(self, malware_name: str, path: str):
		self.malware_name = malware_name
		self.path = path
		self.trace_array = None
		self._get_trace()
	
	def _get_trace(self) -> None:
		pids = list(set(get_malware_pids(self.malware_name, path=self.path)))
		return_array = np.empty((0, 4))
		for pid in pids:
			if not os.path.isfile(f"{self.path}frida_{pid}.txt.gz"):
				# print(f"Trace for {pid} does not exist")
				continue
			header = ["time", "api", "category"]
			frida_str = StringIO(gzip_to_string(f"{self.path}frida_{pid}.txt.gz"))
			dataframe = pd.read_csv(frida_str, names=header)
			dataframe.drop(
				dataframe.loc[dataframe['api'] == 'error'].index, inplace=True)
			dataframe.drop_duplicates(
				subset=['time', 'api'], keep='first', inplace=True)
			dataframe.reset_index(drop=True, inplace=True)
			dataframe["time_int"] = (dataframe["time"] * 1000000).astype(np.int64)
			# To numpy
			np_array_df = dataframe.to_numpy()
			# Stack to the final array
			return_array = np.vstack((return_array, np_array_df))
		self.trace_array = return_array
		print(len(self.trace_array))
	
	def get_merge_trace(self) -> np.ndarray:
		return self.trace_array


def get_segmented_flow_syscalls(segmented_flow: np.ndarray, malware_process_name: str, path: str = "./",
                                time_delay_allowed: int = 0) -> np.ndarray:
	"""
	Get a list of API calls corresponding to the segmented flow
	:param segmented_flow:
	:param malware_process_name:
	:param path:
	:param time_delay_allowed:
	:return:
	"""
	min_time, max_time = segmented_flow[0][2]
	for group in segmented_flow[1:]:
		# ['HANDSHAKE', [0, 2], [1612708961378936, 1612708961422139]]
		timea, timeb = group[2]
		min_time = min(timea, min_time)
		max_time = max(timeb, max_time)
	# print(min_time, max_time, max_time - min_time)
	trace_extractor = MalwareTraceExtractor(malware_name=malware_process_name, path=path)
	calls = trace_extractor.get_merge_trace()
	# calls = get_malware_traces_merged(malware_process_name, path=path)
	returned_calls = np.empty((0, 4))
	for call in calls:
		mapping = pure_time_map(min_time, max_time, call[3], time_delay_allowed)
		if mapping is not None:
			returned_calls = np.vstack((returned_calls, call))
	return returned_calls


if __name__ == "__main__":
	PATH = "trickbot1_1/"
