def test_ascad_iterator(self):
"""
Check whether the AICorrSignalIterator returns the same output as load_ascad
:return:
"""
from ASCAD_train_models import load_ascad
conf = Namespace(
input_type=AIInputType.SIGNAL,
leakage_model=LeakageModelType.SBOX_OH,
max_cache=None,
augment_roll=False,
augment_noise=False,
augment_shuffle=False,
normalize=False,
traces_per_set=50000,
online=False,
dataset_id='test',
format='ascad',
reference_signal=np.array([0] * 700),
actions=[Action('window[0,700]')],
cnn=False,
key_low=2,
key_high=3,
windowing_method='rectangular',
norank=False,
)
ascad_path = "./datasets/ASCAD/ASCAD_data/ASCAD_databases/ASCAD.h5"
iterator = AICorrSignalIterator([ascad_path + "-train"],
conf,
batch_size=256)
x = np.zeros((512, 700))
y = np.zeros((512, 256))
inputs, labels = next(iterator)
x[0:256] = inputs
y[0:256] = labels
inputs, labels = next(iterator)
x[256:512] = inputs
y[256:512] = labels
(X_profiling, Y_profiling), (X_attack,
Y_attack), (meta_profiling,
meta_attack) = load_ascad(
ascad_path, True)
x_ascad = X_profiling
y_ascad = to_categorical(Y_profiling, num_classes=256)
for i in range(0, 512):
self.assertListEqual(list(x[i]), list(x_ascad[i]))
for i in range(0, 512):
self.assertListEqual(list(y[i]), list(y_ascad[i]))
def test_iterator_wrapping(self):
conf = Namespace(
input_type=AIInputType.SIGNAL,
leakage_model=LeakageModelType.SBOX_OH,
max_cache=None,
augment_roll=False,
augment_noise=False,
augment_shuffle=False,
normalize=False,
traces_per_set=32,
online=False,
dataset_id='test',
format='cw',
reference_signal=np.array([0] * 128),
actions=[],
cnn=False,
key_low=2,
key_high=3,
norank=False,
)
iterator = AICorrSignalIterator([
"./datasets/unit-test/test_traces.npy",
"./datasets/unit-test/test2_traces.npy"
],
conf,
batch_size=48)
inputs, labels = next(iterator)
for i in range(0, 48):
self.assertListEqual(list(inputs[i]), [i] * 128)
self.assertListEqual(
list(labels[i]),
list(to_categorical(sbox[1 ^ 0], num_classes=256)))
self.assertEqual(inputs.shape, (48, 128))
self.assertEqual(labels.shape, (48, 256))
inputs, labels = next(iterator)
for i in range(0, 48):
self.assertListEqual(list(inputs[i]), [(i + 48) % 64] * 128)
self.assertEqual(inputs.shape, (48, 128))
self.assertEqual(labels.shape, (48, 256))
def test_corrtrain_correlation_multi(self):
from leakagemodels import LeakageModel
"""
Artificial example to test AICorrNet and trace processing with multiple leakage values and multiple subkeys.
"""
# ------------------------------
# Generate data
# ------------------------------
traces = [ # Contains abs(trace). Shape = [trace, point]
[1, 1, 1, -15],
[-4, 2, 2, -12],
[10, 3, 3, 8],
[8, 1, 1, -14],
[9, 0, -3, 8],
]
plaintexts = [
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 13, 13, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 15, 15, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 8, 8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
]
keys = [
[0, 2, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 2, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 2, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 2, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
[0, 2, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
]
# Convert to numpy
traces = np.array(traces)
plaintexts = np.array(plaintexts)
keys = np.array(keys)
trace_set = TraceSet(name='test',
traces=traces,
plaintexts=plaintexts,
keys=keys)
# ------------------------------
# Preprocess data
# ------------------------------
conf = Namespace(
max_cache=0,
augment_roll=False,
augment_noise=False,
normalize=False,
traces_per_set=4,
online=False,
dataset_id='qa',
cnn=False,
leakage_model=LeakageModelType.AES_MULTI,
input_type=AIInputType.SIGNAL,
augment_shuffle=True,
n_hidden_layers=1,
n_hidden_nodes=256,
activation='leakyrelu',
metric_freq=100,
regularizer=None,
reglambda=0.001,
model_suffix=None,
use_bias=True,
batch_norm=True,
hamming=False,
key_low=1,
key_high=3,
loss_type='correlation',
lr=0.001,
epochs=5000,
batch_size=512,
norank=False,
)
it_dummy = AICorrSignalIterator([],
conf,
batch_size=10000,
request_id=None,
stream_server=None)
x, y = it_dummy._preprocess_trace_set(trace_set)
# ------------------------------
# Train and obtain encodings
# ------------------------------
model = ai.AICorrNet(conf, input_dim=4, name="test")
print(model.info())
rank_cb = rank.CorrRankCallback(conf,
'/tmp/deleteme/',
save_best=False,
save_path=None)
rank_cb.set_trace_set(trace_set)
if model.using_regularization:
print(
"Warning: cant do correlation loss test because regularizer will influence loss function"
)
return
# Find optimal weights
print("The x (EM samples) and y (leakage model values) are:")
print(x)
print(y)
print(
"When feeding x through the model without training, the encodings become:"
)
print(model.predict(x))
print("Training now")
model.train_set(x,
y,
save=False,
epochs=conf.epochs,
extra_callbacks=[rank_cb])
print("Done training")
# Get the encodings of the input data using the same approach used in ops.py corrtest (iterate over rows)
result = []
for i in range(0, x.shape[0]):
result.append(
model.predict(np.array([x[i, :]], dtype=float))[0]
) # Result contains sum of points such that corr with y[key_index] is maximal for all key indices. Shape = [trace, 16]
result = np.array(result)
print(
"When feeding x through the model after training, the encodings for key bytes %d to %d become:\n %s"
% (conf.key_low, conf.key_high, str(result)))
# ------------------------------
# Check loss function
# ------------------------------
# Evaluate the model to get the loss for the encodings
predicted_loss = model.model.evaluate(x, y, verbose=0)
# Manually calculate the loss using numpy to verify that we are learning a correct correlation
calculated_loss = 0
num_keys = (conf.key_high - conf.key_low)
num_outputs = LeakageModel.get_num_outputs(conf) // num_keys
for i in range(0, num_keys):
subkey_hws = y[:, i * num_outputs:(i + 1) * num_outputs]
subkey_encodings = result[:, i * num_outputs:(i + 1) * num_outputs]
print("Subkey %d HWs : %s" % (i + conf.key_low, str(subkey_hws)))
print("Subkey %d encodings: %s" %
(i + conf.key_low, str(subkey_encodings)))
y_key = subkey_hws.reshape([-1, 1])
y_pred = subkey_encodings.reshape([-1, 1])
print("Flattened subkey %d HWs : %s" %
(i + conf.key_low, str(y_key)))
print("Flattened subkey %d encodings: %s" %
(i + conf.key_low, str(y_pred)))
# Calculate correlation (numpy approach)
corr_key_i = np.corrcoef(y_pred[:, 0], y_key[:, 0],
rowvar=False)[1, 0]
print("corr_num: %s" % corr_key_i)
calculated_loss += 1.0 - corr_key_i
print("These values should be close:")
print("Predicted loss: %s" % str(predicted_loss))
print("Calculated loss: %s" % str(calculated_loss))
self.assertAlmostEqual(predicted_loss, calculated_loss, places=2)