#!/usr/bin/env julia

include("util.jl")
include("sbox.jl")
include("plotutils.jl")

import CSV
using Formatting

function parse_csv(filename)
    data = CSV.read(filename, header=0)
    plaintexts::Matrix{UInt8} = convert(Matrix{UInt8}, data[:, 1:16])#[1:50, :]
    timings::Matrix{UInt32} = convert(Matrix{UInt32}, data[:, 17:17])#[1:50]
    return plaintexts, timings
end

function break_encryption()
    plaintexts, timings = parse_csv("timing.csv")
    t_values = Vector{Float64}(undef, 0x100)
    for keybyte=1:16
        for key=0:0xFF
            msb_set = (sbox[(plaintexts[:, keybyte] .⊻ key) .+ 1] .& 0x80) .!= 0
            group_slow = timings[msb_set]
            group_fast = timings[msb_set.==false]
            t_values[key + 1] = t_val(group_fast, group_slow)
        end
        print(format("{:02x} ", argmax(t_values) - 1))
    end
    println()
    plot_discrete_tval(0:0xFF, t_values, "key")
end

function roundkey_g(last, round)
	rc = [0b00000001, 0b00000010, 0b00000100, 0b00001000, 0b00010000, 0b00100000, 0b01000000, 0b10000000, 0x1b, 0x1b << 1]
	next = Vector{UInt8}(undef, 4)
	next[4] = sbox[last[1] + 1]
	next[1] = sbox[last[2] + 1] ⊻ rc[round]
	next[2:3] = sbox[last[3:4] .+ 1]
	return next
end

function build_key_schedule(initial_key)
	round_keys = Matrix{UInt8}(undef, 11, 16)
	round_keys[1, :] = initial_key
	for r=2:11
		round_keys[r, 1:4] = round_keys[r - 1, 1:4] .⊻ roundkey_g(round_keys[r - 1, 13:16], r - 1)
		round_keys[r, 5:8] = round_keys[r, 1:4] .⊻ round_keys[r - 1, 5:8]
		round_keys[r, 9:12] = round_keys[r, 5:8] .⊻ round_keys[r - 1, 9:12]
		round_keys[r, 13:16] = round_keys[r, 9:12] .⊻ round_keys[r - 1, 13:16]
	end
	return round_keys
end

function break_decryption()
    keys, timings = parse_csv("timing.csv")
	round_keys = Array{UInt8, 3}(undef,11, size(keys, 1), 16)
	for keyno=1:size(keys, 1)
		round_keys[:, keyno, :] = build_key_schedule(keys[keyno, :])
	end
    t_values = Vector{Float64}(undef, 0x100)
	r_shiftrows=[0, 13, 10, 7, 4, 1, 14, 11, 8, 5, 2, 15, 12, 9, 6, 3] .+ 1
    for secretbyte=1:16
        for secret=0:0xFF
			# TODO Shift rows must be respected here
			msbs_set = ((round_keys[10, :, r_shiftrows[secretbyte]] .⊻ r_sbox[(secret .⊻ round_keys[11, :, secretbyte]) .+ 1]) .& 0xE0) .!= 0
			#msbs_set = (sbox[(keys[:, secretbyte] .⊻ secret) .+ 1] .& 0x08) .!= 0
            group_slow = timings[msbs_set]
            group_fast = timings[msbs_set.==false]
            t_values[secret + 1] = t_val(group_fast, group_slow)
        end
        print(format("{:02x} ", argmax(t_values) - 1))
        if secretbyte==1
            plt = plot_discrete_tval(0:0xFF, t_values, "secret data")
			display(plt)
        end
    end
    println()
end

break_decryption()