from trojan_insertions.steganography_encode_decode import SteganographyEncodeDecode
from trojan_insertions.steganography_utility import SteganographyUtility
from trojan_insertions.direct_encode_decode import DirectEncodeDecode
import re

class TrojanExperimentUtility():
    def __init__(self, key=60, key2=3, window_size=1, start_marker='#', end_marker='#'):
        self.encode_decode = SteganographyEncodeDecode(key, key2, window_size, start_marker, end_marker)
        self.st_utility = SteganographyUtility(key2=key2)
        self.normal_encode_decode = DirectEncodeDecode(window_size, start_marker, end_marker)
        self.start_marker_size = len(start_marker)
        self.end_marker_size = len(end_marker)
        self.start_marker = start_marker
        self.end_marker = end_marker

    def make_payload(self, payload, wrap_rand_neucleotide=False, normal_encoding=False):
        coverDNA = self.st_utility.get_random_cover_message(payload, wrap_rand_neucleotide=wrap_rand_neucleotide)
        if normal_encoding:
            encodedDNA = self.normal_encode_decode.encode(payload)
        else:
            encodedDNA = self.encode_decode.encode(payload, coverDNA)
        return encodedDNA
    def get_decoded_parts(self, encodedDNA, normal_encoding=False):
        if normal_encoding:
            payload_parts = self.normal_encode_decode.findingPayload(encodedDNA)
        else:
            payload_parts = self.encode_decode.findingPayload(encodedDNA)
        return payload_parts
    def to_parts(self, payload, slice_size, start_marker='#', end_marker='#'):
        parts = []
        k = 0
        for i in range(0, len(payload), slice_size):
            # parts.append(start_marker+str(k)+payload[i:i+slice_size]+end_marker)
            parts.append(start_marker+str(k)+','+payload[i:i+slice_size]+end_marker)
            k += 1
        return parts
    # def combine_parts(self, parts, slice_size):
    #     payload = ['' for i in range(0,len(parts))]
    #     for part in parts:
    #         part_without_marker = part[self.start_marker_size:-self.end_marker_size]
    #         idx = part_without_marker[:-(slice_size)]
    #         main_part = part_without_marker[-(slice_size):]
    #         try:
    #             idx = int(idx)
    #             payload[idx] = main_part
    #         except:
    #             pass
    #     return ''.join(payload)
    def combine_parts(self, parts):
        payload = ['' for i in range(0,len(parts))]
        for part in parts:
            try:
                idx = part.split(',')[0]
                idx = int(idx.split(self.start_marker)[1])
                main_part = part.split(',')[1]
                main_part = main_part.split(self.end_marker)[0]
                payload[idx] = main_part
            except:
                pass
        return ''.join(payload)

    def get_connections(self, payload, regex_pattern = r"[a-zA-Z0-9-]{1,}[.]{1}[a-zA-Z]{1,}[a-zA-Z0-9-]{1,}:[0-9]{1,}"):
        pattern = re.compile(regex_pattern)
        matchedStrings = pattern.findall(payload)
        if len(matchedStrings) > 0:
            return True, matchedStrings
        return False, []




