formats.py

from abc import ABC, abstractmethod
import base64
import logging
from math import log2, ceil
from subprocess import run


logging.basicConfig(level=logging.INFO)

zero_byte = 0x0.to_bytes(1, byteorder="big")


def xor(param1, param2):
    return bytes((x ^ y) for (x, y) in zip(param1, param2))


class Message(ABC):

    # AES block size
    block_size = 16

    @abstractmethod
    def get_block_amount(self):
        pass

    @abstractmethod
    def get_ciphertext_block(self, nr):
        pass

    @abstractmethod
    def insert_in_ciphertext(self, block_nr, *content_vec):
        pass

    @abstractmethod
    def adapt_length(self):
        pass

    @staticmethod
    def calculate_needed_length_bytes(new_length):
        needed_bits = int(log2(new_length)) + 1
        needed_bytes = ceil(needed_bits / 8)
        return needed_bytes


class P7m(Message):

    def __init__(self, p7m, iv_offset, ciphertext_offset, ciphertext_length, length_places):

        self.msg = p7m
        self.msg_bytes = bytearray(base64.b64decode(self.msg))

        self.iv_start = iv_offset
        self.ciphertext_offset = ciphertext_offset
        self.ciphertext_length = ciphertext_length
        self.length_places = length_places
        self.length_diff = 0

    def get_block_amount(self):
        return int(self.get_ciphertext_length() / self.block_size)

    def get_ciphertext_length(self):
        return len(self.msg_bytes[self.ciphertext_offset:])

    def get_iv(self):
        iv_bytes = bytes(self.msg_bytes[self.iv_start: self.iv_start + self.block_size])
        return iv_bytes

    def get_ciphertext_block(self, nr):
        ciphertext = self.msg_bytes[self.ciphertext_offset:]
        return ciphertext[(nr-1) * self.block_size: nr * self.block_size]

    def insert_in_ciphertext(self, block_nr, *contentv):
        # Determine the place to insert
        place = self.ciphertext_offset + block_nr * self.block_size

        for content in reversed(contentv):
            self.msg_bytes[place:place] = content

        self.length_diff = len(contentv) * self.block_size
        self.adapt_length()

    def adapt_length(self):

        for i in self.length_places:
            # For each element, which has the ciphertext nested

            first_len_byte = self.msg_bytes[i]

            if first_len_byte > 0b10000000:  # long form

                current_amount_length_bytes = first_len_byte - 0x80
                current_length_bytes = self.msg_bytes[i+1: i+1+current_amount_length_bytes]
                current_length = int.from_bytes(current_length_bytes, byteorder="big")
                new_length = current_length + self.length_diff
                needed_bytes = self.calculate_needed_length_bytes(new_length)

                if needed_bytes > current_amount_length_bytes:  # Add new byte(s)

                    diff = needed_bytes - current_amount_length_bytes
                    self.msg_bytes[i+1: i+1] = zero_byte * diff
                    self.msg_bytes[i+1: i+1+needed_bytes] = new_length.to_bytes(needed_bytes, byteorder="big")

                    self.msg_bytes[i] += diff
                    self.length_diff += diff
                    self.ciphertext_offset += diff

                    if i != self.length_places[0]:
                        self.length_places[0] += diff

                    logging.info("Added length {} byte(s) after byte {}".format(needed_bytes, i))
                else:
                    start = i + 1
                    end = i + 1 + current_amount_length_bytes
                    self.msg_bytes[start: end] = new_length.to_bytes(needed_bytes, byteorder="big")

            else:  # short form
                new_length = first_len_byte + self.length_diff

                if new_length >= 0x80:  # switch to long form

                    needed_bytes = self.calculate_needed_length_bytes(new_length)
                    self.msg_bytes[i] = 0x80 + needed_bytes
                    logging.info("Switched to long form at byte {}".format(i))

                    self.msg_bytes[i+1: i+1] = new_length.to_bytes(needed_bytes, byteorder="big")
                    self.length_diff += needed_bytes
                    self.ciphertext_offset += needed_bytes
                    logging.info("Added {} length byte(s) after byte {}".format(needed_bytes, i))

                else:
                    self.msg_bytes[i] += self.length_diff

    def format_properly(self):

        # Convert bytes to base64 string
        b64_encoded_bytes = base64.b64encode(self.msg_bytes)
        msg_b64_string = str(b64_encoded_bytes, "ascii")

        # Insert line breaks as recommended
        formatted = "\n".join(msg_b64_string[pos: pos + 64] for pos in range(0, len(msg_b64_string), 64))

        return formatted


class OpenPgpMsg(Message):

    def __init__(self, bin_msg, pkesk_len, seipd_hlen):

        self.data = bin_msg
        self.pkesk_len = pkesk_len

        self.seipd_offset = self.pkesk_len
        self.seipd_hlen = seipd_hlen

        self.mdc_hlen = 2
        self.mdc_plen = 20
        self.mdc_len = self.mdc_hlen + self.mdc_plen

    def get_seidp_plen(self):
        return len(self.data[self.seipd_offset + self.seipd_hlen:])

    def get_seidp_len(self):
        return self.seipd_hlen + self.get_seidp_plen()

    def get_ld_plen(self):
        bytes_before_ld = self.block_size + 5
        return self.get_seidp_plen() - bytes_before_ld - self.mdc_len

    def get_seipd_body_offset(self):
        return self.seipd_offset + self.seipd_hlen + 1

    def get_ciphertext_block(self, nr):
        off = self.get_seipd_body_offset()
        return self.data[off + (nr-1)*self.block_size: off + nr*self.block_size]

    def get_block_amount(self):
        return int(self.get_seidp_plen() / self.block_size)

    def insert_in_ciphertext(self, block_nr, *content_vec):
        place = self.get_seipd_body_offset() + self.block_size * block_nr

        for content in reversed(content_vec):
            self.data[place:place] = content

        self.adapt_length()

    @staticmethod
    def ctb_is_in_new_format(ctb):
        return ctb > 0b11000000

    @staticmethod
    def determine_length_bytes_amount(first_byte):
        if first_byte in range(192):
            return 1
        elif first_byte in range(192, 223):
            return 2
        else:
            raise NotImplementedError

    @staticmethod
    def encode_len(num):
        tmp1 = num - 192
        mask = 0b11111111

        second = tmp1 & mask
        tmp2 = tmp1 - second

        tmp3 = tmp2 >> 8
        first = tmp3 + 192

        res1 = first.to_bytes(1, byteorder="big")
        res2 = second.to_bytes(1, byteorder="big")

        return res1 + res2

    def create_new_header_block(self):
        quick_check_bytes = 2 * zero_byte

        # CTB for Tag 11 (Literal Data) in new format
        ctb = 0b11001011.to_bytes(1, byteorder="big")

        # New length
        new_len = self.get_ld_plen() + (2 * self.block_size)

        # Determine if one or two bytes are needed and then store the length.
        if new_len < 192:
            # One byte
            new_plen_byte = new_len.to_bytes(1, byteorder="big")
            remaining_bytes = b'Conten'

        elif new_len < 8383:
            # Two bytes
            new_plen_byte = OpenPgpMsg.encode_len(new_len - 1)
            self.adapt_length()
            remaining_bytes = b'Conte'

        else:
            # more than 8383 bytes
            raise NotImplementedError

        mode = 0x62.to_bytes(1, byteorder="big")

        # Include next block of random bytes
        name_len = 0x1f.to_bytes(1, byteorder="big")

        date = 4 * zero_byte

        return quick_check_bytes + ctb + new_plen_byte + mode + name_len + date + remaining_bytes

    def adapt_length(self):

        ctb = self.data[self.seipd_offset]

        if self.ctb_is_in_new_format(ctb):

            first_len_byte = self.data[self.seipd_offset + 1]
            current_addl_length_bytes = self.determine_length_bytes_amount(first_len_byte)
            needed_bytes = 1 if self.get_seidp_plen() < 192 else 2

            if current_addl_length_bytes == needed_bytes:
                off = self.seipd_offset + 1
                if current_addl_length_bytes == 1:
                    self.data[off] = self.get_seidp_plen()
                elif current_addl_length_bytes == 2:
                    self.data[off: off + 2] = self.encode_len(self.get_seidp_plen())
                else:
                    raise NotImplementedError

            elif current_addl_length_bytes < needed_bytes:
                if current_addl_length_bytes == 1:
                    len_bytes = self.encode_len(self.get_seidp_plen())
                    offset = self.seipd_offset + 2
                    self.data[offset: offset] = b'0'
                    self.data[offset - 1: offset + 1] = len_bytes
                    self.seipd_hlen += 1
                else:
                    raise NotImplementedError

        else:
            # determine value of two least significant bits
            bit_mask = 0b11
            res = ctb & bit_mask

            # determine current length settings
            current_addl_length_bytes = 2**res
            current_length = self.data[self.seipd_offset + 1: self.seipd_offset + current_addl_length_bytes]
            current_length_int = int.from_bytes(current_length, byteorder="big")

            # determine new length settings and adapt
            new_length = current_length_int + self.block_size
            needed_bytes = self.calculate_needed_length_bytes(self.get_seidp_plen())

            if current_addl_length_bytes < needed_bytes:
                raise NotImplementedError
            else:
                start = self.seipd_offset + 1
                end = self.seipd_offset + current_addl_length_bytes
                self.data[start: end] += new_length.to_bytes(1, byteorder="big")

    @staticmethod
    def enarmor():
        run(["gpg", "--batch", "--yes", "--enarmor", "modified.eml.gpg"], cwd="ciphertext_files/openpgp/")
        with open("ciphertext_files/openpgp/modified.eml.gpg.asc", "r") as f:
            read = f.read()
        f.closed
        replaced = read.replace("ARMORED FILE", "MESSAGE")
        return replaced