#!/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])
    timings::Matrix{UInt32} = convert(Matrix{UInt32}, data[:, 17:17])
    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 gf_mult2(value::UInt8)::UInt8
	result = value << 1
	if value & 0x80 != 0
		result ⊻= value
	end
	return result
end

function gf_mult3(value::UInt8)::UInt8
	return gf_mult2(value) ⊻ value
end

function break_decryption()
    keys, timings = parse_csv("timing.csv")
	round_keys = Array{UInt8, 3}(undef,11, size(keys, 1), 16)
	@Threads.threads for keyno=1:size(keys, 1)
		round_keys[:, keyno, :] = build_key_schedule(keys[keyno, :])
	end
    t_values = Vector{Float64}(undef, 0x100)
	shiftrows = [0, 5, 10, 15, 4, 9, 14, 3, 8, 13, 2, 7, 12, 1, 6, 11] .+ 1
	r_shiftrows = [0, 13, 10, 7, 4, 1, 14, 11, 8, 5, 2, 15, 12, 9, 6, 3] .+ 1
	byte1 = shiftrows[1]
	byte2 = shiftrows[2]
	byte3 = shiftrows[3]
	byte4 = shiftrows[4]
	high_tvalues = [Vector{Any}(undef, 0) for i=1:Threads.nthreads()]
	@Threads.threads for p1=0xA0:0xA0
		for p2=0x00:0x00, p3=0x70:0x80, p4=0xD0:0xE0
			msbs_set = (gf_mult2.(sbox[p1 .⊻ round_keys[1, :, byte1] .+ 1]) .⊻ gf_mult3.(sbox[p2 .⊻ round_keys[1, :, byte2] .+ 1]) .⊻ sbox[p3 .⊻ round_keys[1, :, byte3] .+ 1] .⊻ sbox[p4 .⊻ round_keys[1, :, byte4] .+ 1]) .& 0xC0 .!= 0
			group_slow = timings[msbs_set]
			group_fast = timings[msbs_set.==false]
			t_value = t_val(group_fast, group_slow)
			if abs(t_value) > 4
				push!(high_tvalues[Threads.threadid()], (p1, p2, p3, p4, t_value))
			end
		end
	end
	println(high_tvalues)
end

@time break_decryption()