from common.common import *

class ExploitsBuilder(object):

	def __init__(self, test_cases, config):
		self.config = config
		self.case_id = config['case_id'].decode("utf-8")
		self.mode = config['mode']
		self.test_cases = self.fixup_test_case_data(test_cases)


	def fixup_test_case_data(self, t):
		if self.config['mode'] == 'c':
			t = recursive_fixup(t, b"attacker@example.com", self.config["client_mode"]["username"])
			t = recursive_fixup(t, b"admin@example.com", self.config["legitimate_site_address"])
			t= recursive_fixup(t, b"victim@victim.com", self.config["victim_address"])
		else:
			t = recursive_fixup(t, b"attack.com", self.config["attacker_site"])
			t = recursive_fixup(t, b"admin@legitimate.com", self.config["legitimate_site_address"])
			legitimate_site = self.config["legitimate_site_address"].split(b"@")[1]
			t = recursive_fixup(t, b"legitimate.com", legitimate_site)
			t= recursive_fixup(t, b"victim@victim.com", self.config["victim_address"])

		if self.config['subject_header']:
			t[self.case_id]['data']['subject_header'] = self.config['subject_header']
		if self.config['to_header']:
			t[self.case_id]['data']['to_header'] = self.config['to_header']
		if self.config['body']:
			t[self.case_id]['data']['body'] =self.config['body']

		return t

	def generate_msg_data(self):
		test_cases = self.test_cases
		case_id = self.case_id
		msg_content = test_cases[case_id]["data"]
		dkim_para = test_cases[case_id].get("dkim_para")
		if dkim_para != None:
			dkim_msg =   dkim_para["sign_header"] +b"\r\n\r\n" + msg_content["body"]
			dkim_header = generate_dkim_header(dkim_msg, dkim_para)
			msg = msg_content["from_header"] + dkim_header + msg_content["to_header"] + msg_content["subject_header"] + msg_content["other_headers"] + msg_content["body"]
		else:
			msg = msg_content["from_header"] + msg_content["to_header"] + msg_content["subject_header"] + msg_content["other_headers"] + msg_content["body"]
		return msg

	def generate_smtp_seqs(self):
		test_cases = self.test_cases
		case_id = self.case_id

		smtp_seqs = {
			"helo": test_cases[case_id]["helo"],
			"mailfrom": test_cases[case_id]["mailfrom"],
			"rcptto": test_cases[case_id]["rcptto"],
			"msg_content": self.generate_msg_data()
		}
		return smtp_seqs
