aiinputs.py

# ----------------------------------------------------
# Electromagnetic Mining Array (EMMA)
# Copyright 2018, Pieter Robyns
# ----------------------------------------------------

import numpy as np
from emutils import EMMAException, int_to_one_hot
from leakagemodels import LeakageModel


class AIInputType:
    """
    Class that defines all possible types of inputs for the ML models. Input classes must have an attribute
    'input_type' with one of the values defined in this class.
    """
    SIGNAL = 'signal'
    SIGNAL_PLAINTEXT = 'signal_plaintext'

    # For testing purposes
    SIGNAL_KEY = 'signal_key'
    SIGNAL_PLAINTEXT_KEY = 'signal_plaintext_key'
    PLAINTEXT_KEY = 'plaintext_key'
    PLAINTEXT_KEY_OH = 'plaintext_key_oh'
    SIGNAL_LEAKAGE = 'signal_leakage'
    RANDOM = 'random'

    @classmethod
    def choices(cls):
        """
        Get all possible AIInputTypes in list form
        :return:
        """
        c = []
        for k, v in cls.__dict__.items():
            if k[:2] != '__' and type(v) is str:
                c.append(v)
        return c


class AIInputMeta(type):
    """
    Metaclass used for checking whether the child class contains a valid input_type attribute.
    """
    class BadAIInputClassException(EMMAException):
        pass

    class InvalidInputTypeException(EMMAException):
        pass

    def __new__(mcs, name, bases, class_dict):
        if bases != (object,):  # Do not validate LeakageModel class
            if 'input_type' not in class_dict:
                raise AIInputMeta.BadAIInputClassException
            if class_dict['input_type'] not in AIInputType.choices():
                raise AIInputMeta.InvalidInputTypeException
        return type.__new__(mcs, name, bases, class_dict)


class AIInput(object, metaclass=AIInputMeta):
    """
    AI input base class.
    """
    class UnknownAIInputException(EMMAException):
        pass

    def __new__(cls, conf):
        """
        Called when instantiating an AIInput object. Returns an instance of the appropriate class depending on the
        input_type parameter.
        :param conf:
        :return:
        """
        for subclass in cls._get_subclasses():
            if subclass.input_type == conf.input_type:
                return object.__new__(subclass)  # Avoid recursion by calling object.__new__ instead of cls.__new__
        raise AIInput.UnknownAIInputException

    def __init__(self, conf):
        self.conf = conf

    @classmethod
    def _get_subclasses(cls):
        for subclass in cls.__subclasses__():
            if cls is not object:
                for subsubclass in subclass._get_subclasses():  # Also yield children of children
                    yield subsubclass
            yield subclass

    def get_trace_inputs(self, trace):
        raise NotImplementedError

    def get_trace_set_inputs(self, trace_set):
        """
        Givem a trace set, returns inputs suitable for training an AI model.
        :param trace_set:
        :return:
        """
        inputs = []

        for trace in trace_set.traces:
            inputs.append(self.get_trace_inputs(trace))

        result = np.array(inputs)

        # CNNs expect a channels dimension
        if self.conf.cnn:
            result = np.expand_dims(result, axis=-1)

        return result


class SignalAIInput(AIInput):
    input_type = AIInputType.SIGNAL

    def get_trace_inputs(self, trace):
        return trace.signal


class SignalPlaintextAIInput(AIInput):
    input_type = AIInputType.SIGNAL_PLAINTEXT

    def get_trace_inputs(self, trace):
        return np.concatenate((trace.signal, trace.plaintext))


class SignalKeyAIInput(AIInput):
    input_type = AIInputType.SIGNAL_KEY

    def get_trace_inputs(self, trace):
        return np.concatenate((trace.signal, trace.key))


class SignalPlaintextKeyAIInput(AIInput):
    input_type = AIInputType.SIGNAL_PLAINTEXT_KEY

    def get_trace_inputs(self, trace):
        return np.concatenate((trace.signal, trace.plaintext, trace.key))


class PlaintextKeyAIInput(AIInput):
    input_type = AIInputType.PLAINTEXT_KEY

    def get_trace_inputs(self, trace):
        return np.concatenate((trace.plaintext, trace.key))


class PlaintextKeyOHAIInput(AIInput):
    input_type = AIInputType.PLAINTEXT_KEY_OH

    def get_trace_inputs(self, trace):
        result = []
        for p in trace.plaintext:
            result.append(int_to_one_hot(p, 256))
        for k in trace.key:
            result.append(int_to_one_hot(k, 256))
        return np.concatenate(result)


class SignalLeakageAIInput(AIInput):
    input_type = AIInputType.SIGNAL_LEAKAGE

    def __init__(self, conf):
        super().__init__(conf)
        self.leakage_model = LeakageModel(conf)

    def get_trace_inputs(self, trace):
        leakages = []

        for k in range(16):
            leakage = self.leakage_model.get_trace_leakages(trace, k)
            if isinstance(leakage, list) or isinstance(leakage, np.ndarray):
                leakages.extend(list(leakage))
            else:
                leakages.append(leakage)
        leakages = np.array(leakages)

        return np.concatenate((trace.signal, leakages))


class RandomInput(AIInput):
    input_type = AIInputType.RANDOM

    def get_trace_inputs(self, trace):
        return np.random.uniform(0.0, 1.0, len(trace.signal))