#!ruby
## test base block features
require "minitest/autorun"

require_relative "../cryptosystems/deoxys.rb"
require_relative "../writers/graphviz.rb"

def rows_p()
    return [0, 1, 2, 3, 5, 6, 7, 4, 10, 11, 8, 9, 15, 12, 13, 14]
end

class TestDeoxys_stepbystep < Minitest::Unit::TestCase
    def setup()
        @m = [
            0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00
        ]
        @tk1 = [ #TK1
            0x10, 0x11, 0x12, 0x13,
            0x14, 0x15, 0x16, 0x17,
            0x18, 0x19, 0x1a, 0x1b,
            0x1c, 0x1d, 0x1e, 0x1f 
        ] 
        @tk2 = [ #TK2
            0x12, 0x02, 0x12, 0x22,
            0x32, 0x42, 0x52, 0x62,
            0x70, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00
        ]
        @R1 = [ # C
            0x6a, 0x59, 0xf1, 0x38,
            0x33, 0x5c, 0x67, 0xed,
            0x03, 0x99, 0xc3, 0xe4,
            0x54, 0x9b, 0x46, 0x47
        ]
        @prewhite = [
            0x03, 0x11, 0x04, 0x39,
            0x09, 0x78, 0x6b, 0x5a,
            0x68, 0x19, 0x1a, 0x1b,
            0x1c, 0x1d, 0x1e, 0x1f
        ]
        @k3 = [
            0x00, 0x00, 0x00, 0x01,
            0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00
        ]
        @permut = [
            0x00, 0x04, 0x08, 0x0c,
            0x01, 0x05, 0x09, 0x0d,
            0x02, 0x06, 0x0a, 0x0e,
            0x03, 0x07, 0x0b, 0x0f
        ]
        @permut1 = [
            0x00, 0x01, 0x02, 0x03,
            0x04, 0x05, 0x06, 0x07,
            0x08, 0x09, 0x0a, 0x0b,
            0x0c, 0x0d, 0x0e, 0x0f
        ]
        @al = [
            0x01, 0x05, 0x09, 0x0d,
            0x06, 0x0a, 0x0e, 0x02,
            0x0b, 0x0f, 0x03, 0x07,
            0x0c, 0x00, 0x04, 0x08
        ]
    end
    def test_d0() #key set-up
        dag = Dag.new([],[],[],[])
        x,_,_ = dag.register_block(*input_block("X", BYTE_RANGE, [4,4]))
        z = 256
        if z == 256
            k2,_,_ = dag.register_block(*input_block("TK1", BYTE_RANGE, [4,4])) #entrada de tk1
            k1,_,_ = dag.register_block(*input_block("TK2", BYTE_RANGE, [4,4])) #entrada de tk2
            dag.set_inputs(x.flatten+k2.flatten+k1.flatten)
        elsif z == 384
            k3,_,_ = dag.register_block(*input_block("TK1", BYTE_RANGE, [4,4]))
            k2,_,_ = dag.register_block(*input_block("TK2", BYTE_RANGE, [4,4]))
            k1,_,_ = dag.register_block(*input_block("TK3", BYTE_RANGE, [4,4]))
            dag.set_inputs(x.flatten+k3.flatten+k2.flatten+k1.flatten)
            b3 = k3
            tk3_params = [[7,1],[0],[1],[2],[3],[4],[5],[6]]
        else
            raise "create_skinny_dag: invalid z parameter (sould be in {256,384}, current value: #{z})"
        end
        a = x
        b1 = k1
        b2 = k2
        tk2_params = [[1],[2],[3],[4],[5],[6],[7],[0,2]]
        nb_rounds = 1
        ###################################################################################################
        dag.set_outputs(a.flatten+b1.flatten+b2.flatten)
        assert_equal(dag.run(@m+@tk1+@tk2), [
            0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00,
            0x12, 0x02, 0x12, 0x22,
            0x32, 0x42, 0x52, 0x62,
            0x70, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00,
            0x10, 0x11, 0x12, 0x13,
            0x14, 0x15, 0x16, 0x17,
            0x18, 0x19, 0x1a, 0x1b,
            0x1c, 0x1d, 0x1e, 0x1f 
        ])
    end
    def test_d1() #key set-up
        dag = Dag.new([],[],[],[])
        x,_,_ = dag.register_block(*input_block("X", BYTE_RANGE, [4,4]))
        z = 256
        if z == 256
            k2,_,_ = dag.register_block(*input_block("TK1", BYTE_RANGE, [4,4]))
            k1,_,_ = dag.register_block(*input_block("TK2", BYTE_RANGE, [4,4]))
            dag.set_inputs(x.flatten+k2.flatten+k1.flatten)
        elsif z == 384
            k3,_,_ = dag.register_block(*input_block("TK1", BYTE_RANGE, [4,4]))
            k2,_,_ = dag.register_block(*input_block("TK2", BYTE_RANGE, [4,4]))
            k1,_,_ = dag.register_block(*input_block("TK3", BYTE_RANGE, [4,4]))
            dag.set_inputs(x.flatten+k3.flatten+k2.flatten+k1.flatten)
            b3 = k3
            tk3_params = [[7,1],[0],[1],[2],[3],[4],[5],[6]]
        else
            raise "create_skinny_dag: invalid z parameter (sould be in {256,384}, current value: #{z})"
        end
        a = x
        b1 = k1
        b2 = k2
        tk2_params = [[1],[2],[3],[4],[5],[6],[7],[0,2]]
        nb_rounds = 1
        ###################################################################################################
        nb_rounds.times do |round_number|
            ####### define cryptosystem internals (X message)
            rc,_,_ = dag.register_block(*input_block("RC_#{round_number}", BYTE_RANGE, [4,4], [
                [0x01, 0x02, 0x04, 0x08],
                [rcon[round_number], rcon[round_number], rcon[round_number], rcon[round_number]],
                [0x00, 0x00, 0x00, 0x00],
                [0x00, 0x00, 0x00, 0x00]
            ]))
            # tweakey xor
            if z == 256
                a,_,_ = dag.register_block(*xor_block("ART_#{round_number}", [a,b1,b2,rc]))
            elsif z == 384
                a,_,_ = dag.register_block(*xor_block("ART_#{round_number}", [a,b1,b2,b3,rc])) 
            else
                raise "create_skinny_dag: invalid z parameter (sould be in {256,384}, current value: #{z})"
            end
        end
        dag.set_outputs(a.flatten)
        assert_equal(dag.run(@m+@tk1+@tk2), [
            0x03, 0x11, 0x04, 0x39,
            0x09, 0x78, 0x6b, 0x5a,
            0x68, 0x19, 0x1a, 0x1b,
            0x1c, 0x1d, 0x1e, 0x1f
        ])
    end
    def test_d2() #key update
        dag = Dag.new([],[],[],[])
        x,_,_ = dag.register_block(*input_block("X", BYTE_RANGE, [4,4]))
        z = 256
        if z == 256
            k2,_,_ = dag.register_block(*input_block("TK1", BYTE_RANGE, [4,4]))
            k1,_,_ = dag.register_block(*input_block("TK2", BYTE_RANGE, [4,4]))
            dag.set_inputs(x.flatten+k2.flatten+k1.flatten)
        elsif z == 384
            k3,_,_ = dag.register_block(*input_block("TK1", BYTE_RANGE, [4,4]))
            k2,_,_ = dag.register_block(*input_block("TK2", BYTE_RANGE, [4,4]))
            k1,_,_ = dag.register_block(*input_block("TK3", BYTE_RANGE, [4,4]))
            dag.set_inputs(x.flatten+k3.flatten+k2.flatten+k1.flatten)
            b3 = k3
            tk3_params = [[7,1],[0],[1],[2],[3],[4],[5],[6]]
        else
            raise "create_skinny_dag: invalid z parameter (sould be in {256,384}, current value: #{z})"
        end
        a = x
        b1 = k1
        b2 = k2
        tk2_params = [[1],[2],[3],[4],[5],[6],[7],[0,2]]
        nb_rounds = 14
        round_tweakey = a
        ##################################################################################
        nb_rounds.times do |round_number|
            ########### key update
            b1,_,_ = dag.register_block(*permutation_block("KP_#{round_number}", b1, tweakey_p))
            b2,_,_ = dag.register_block(*lfsr_block("L_#{round_number}", b2, [
                [1,1,1,1],
                [1,1,1,1],
                [1,1,1,1],
                [1,1,1,1]
            ], tk2_params))
            b2,_,_ = dag.register_block(*permutation_block("KP_#{round_number}", b2, tweakey_p))
        end
        rc,_,_ = dag.register_block(*input_block("RC_#{nb_rounds}", BYTE_RANGE, [4,4], [
                [0x01, 0x02, 0x04, 0x08],    
                [rcon[nb_rounds], rcon[nb_rounds], rcon[nb_rounds], rcon[nb_rounds]],
                [0x00, 0x00, 0x00, 0x00],
                [0x00, 0x00, 0x00, 0x00]
            ]))
        # tweakey xor
        a,_,_ = dag.register_block(*xor_block("ART_#{nb_rounds}", [b1,b2,rc]))
        dag.set_outputs(a.flatten)
        assert_equal(dag.run(@m+@tk1+@tk2), [
            0xff, 0xf8, 0x72, 0x5a,
            0x4e, 0x3a, 0xc6, 0xb2,
            0x8e, 0xba, 0x66, 0x32,
            0xef, 0xab, 0x17, 0x03
        ])
    end
    # 0x09, 0x10, 0x18, 0x0e, XOR 0x01, 0x02, 0x04, 0x08, = 0x08, 0x12, 0x1c, 0x06, /2 -> 0x04, 0x09, 0x0e, 0x03,
    # 0x54, 0x4a, 0x40, 0x5e, XOR 0x5e, 0x5e, 0x5e, 0x5e, = 0x0a, 0x14, 0x1e, 0x00, /2 -> 0x05, 0x0a, 0x0f, 0x00,
    # 0x0c, 0x16, 0x18, 0x02, XOR 0x00, 0x00, 0x00, 0x00, = 0x0c, 0x16, 0x18, 0x02, /2 -> 0x06, 0x0b, 0x0c, 0x01,
    # 0x0e, 0x10, 0x1a, 0x04  XOR 0x00, 0x00, 0x00, 0x00  = 0x0e, 0x10, 0x1a, 0x04  /2 -> 0x07, 0x08, 0x0d, 0x02 
    #
    # 0x05, 0x0b, 0x0a, 0x0b, + 01 02 04 08 = 4 9 
    # 0x5b, 0x54, 0x51, 0x5e, + 5e 5e 5e 5e = 5 a 
    # 0x06, 0x0b, 0x0c, 0x01, + 00 00 00 00 = 6 b 
    # 0x07, 0x08, 0x0d, 0x02  + 00 00 00 00 = 7 8 
    def test_permut() #key update
        dag = Dag.new([],[],[],[])
        x,_,_ = dag.register_block(*input_block("X", BYTE_RANGE, [4,4]))
        dag.set_inputs(x.flatten)
        round_number = 1
        a = x
        a,_,_ = dag.register_block(*permutation_block("KP_#{round_number}", a, trans_p))
        dag.set_outputs(a.flatten)
        assert_equal(dag.run(@al), [
            0x01, 0x06, 0x0b, 0x0c,
            0x05, 0x0a, 0x0f, 0x00,
            0x09, 0x0e, 0x03, 0x04,
            0x0d, 0x02, 0x07, 0x08
        ])
    end
    def test_d3()
        dag = Dag.new([],[],[],[])
        z = 256
        x,_,_ = dag.register_block(*input_block("X", BYTE_RANGE, [4,4]))
        if z == 256
            k2,_,_ = dag.register_block(*input_block("TK1", BYTE_RANGE, [4,4]))
            k1,_,_ = dag.register_block(*input_block("TK2", BYTE_RANGE, [4,4]))
            dag.set_inputs(x.flatten+k2.flatten+k1.flatten)
        elsif z == 384
            k3,_,_ = dag.register_block(*input_block("TK1", BYTE_RANGE, [4,4]))
            k2,_,_ = dag.register_block(*input_block("TK2", BYTE_RANGE, [4,4]))
            k1,_,_ = dag.register_block(*input_block("TK3", BYTE_RANGE, [4,4]))
            dag.set_inputs(x.flatten+k3.flatten+k2.flatten+k1.flatten)
            b3 = k3
            tk3_params = [[7,1],[0],[1],[2],[3],[4],[5],[6]]
        else
            raise "create_skinny_dag: invalid z parameter (sould be in {256,384}, current value: #{z})"
        end
        a = x
        b1 = k1
        b2 = k2
        tk2_params = [[1],[2],[3],[4],[5],[6],[7],[0,2]]
        nb_rounds = 1
        nb_rounds.times do |round_number|
            ####### define cryptosystem internals (X message)
            rc,_,_ = dag.register_block(*input_block("RC_#{round_number}", BYTE_RANGE, [4,4], [
                [0x01, 0x02, 0x04, 0x08],    
                [rcon[round_number], rcon[round_number], rcon[round_number], rcon[round_number]],
                [0x00, 0x00, 0x00, 0x00],
                [0x00, 0x00, 0x00, 0x00]
            ]))
            # tweakey xor
            if z == 256
                a,_,_ = dag.register_block(*xor_block("ART_#{round_number}", [a,b1,b2,rc])) # 1 2 4 8 2f 2f 2f 2f 0 0 0 0 0 0 0 0
            elsif z == 384
                a,_,_ = dag.register_block(*xor_block("ART_#{round_number}", [a,b1,b2,b3,rc])) 
            else
                raise "create_skinny_dag: invalid z parameter (sould be in {256,384}, current value: #{z})"
            end
            #substitution S8()
            a,_,_ = dag.register_block(*subcell_block("S_#{round_number}", a, sbox)) # 0x7C 0x77 0xF2 0x30 0x15 0x15 0x15 0x15 0x63 0x63 0x63 0x63 0x63 0x63 0x63 0x63
            #ShiftRows
            a,_,_ = dag.register_block(*permutation_block("SR_#{round_number}", a,rows_p)) # 0x7C 0x15 0x63 0x63 0x15 0x63 0x63 0x30 0x63 0x63 0xf2 0x15 0x63 0x77 0x15 0x63
            # 0 4 8 c -> 0 4 8 c = 0 5 a f
            # 1 5 9 d -> 5 9 d 1 = 4 9 e 3
            # 2 6 a e -> a e 2 6 = 8 d 2 7
            # 3 7 b f -> f 3 7 b = c 1 6 b
            #MixBytes
        end
        dag.set_outputs(a.flatten)
        assert_equal(dag.run(@m+@m+@m), [
            0x7C, 0x15, 0x63, 0x63,
            0x15, 0x63, 0x63, 0x30,
            0x63, 0x63, 0xf2, 0x15,
            0x63, 0x77, 0x15, 0x63
        ])
    end
    def test_shift_rows()
        dag = Dag.new([],[],[],[])
        x,_,_ = dag.register_block(*input_block("X", BYTE_RANGE, [4,4]))
        dag.set_inputs(x.flatten)
        a = x
        a,_,_ = dag.register_block(*permutation_block("K", a,rows_p))
        dag.set_outputs(a.flatten)
        assert_equal(dag.run(@permut), [
            0x0, 0x5, 0xa, 0xf,
            0x1, 0x6, 0xb, 0xc,
            0x2, 0x7, 0x8, 0xd,
            0x3, 0x4, 0x9, 0xe
        ])
    end
    def test_d4()
        dag = Dag.new([],[],[],[])
        z = 256
        x,_,_ = dag.register_block(*input_block("X", BYTE_RANGE, [4,4]))
        if z == 256
            k2,_,_ = dag.register_block(*input_block("TK1", BYTE_RANGE, [4,4]))
            k1,_,_ = dag.register_block(*input_block("TK2", BYTE_RANGE, [4,4]))
            dag.set_inputs(x.flatten+k2.flatten+k1.flatten)
        elsif z == 384
            k3,_,_ = dag.register_block(*input_block("TK1", BYTE_RANGE, [4,4]))
            k2,_,_ = dag.register_block(*input_block("TK2", BYTE_RANGE, [4,4]))
            k1,_,_ = dag.register_block(*input_block("TK3", BYTE_RANGE, [4,4]))
            dag.set_inputs(x.flatten+k3.flatten+k2.flatten+k1.flatten)
            b3 = k3
            tk3_params = [[7,1],[0],[1],[2],[3],[4],[5],[6]]
        else
            raise "create_skinny_dag: invalid z parameter (sould be in {256,384}, current value: #{z})"
        end
        a = x
        b1 = k1
        b2 = k2
        tk2_params = [[1],[2],[3],[4],[5],[6],[7],[0,2]]
        nb_rounds = 1
        nb_rounds.times do |round_number|
            ####### define cryptosystem internals (X message)
            rc,_,_ = dag.register_block(*input_block("RC_#{round_number}", BYTE_RANGE, [4,4], [
                [0x01, 0x02, 0x04, 0x08],    
                [rcon[round_number], rcon[round_number], rcon[round_number], rcon[round_number]],
                [0x00, 0x00, 0x00, 0x00],
                [0x00, 0x00, 0x00, 0x00]
            ]))
            # tweakey xor
            if z == 256
                a,_,_ = dag.register_block(*xor_block("ART_#{round_number}", [a,b1,b2,rc])) # 1 2 4 8 2f 2f 2f 2f 0 0 0 0 0 0 0 0
            elsif z == 384
                a,_,_ = dag.register_block(*xor_block("ART_#{round_number}", [a,b1,b2,b3,rc])) 
            else
                raise "create_skinny_dag: invalid z parameter (sould be in {256,384}, current value: #{z})"
            end
            #substitution S8()
            a,_,_ = dag.register_block(*subcell_block("S_#{round_number}", a, sbox))       # 0x7C 0x77 0xF2 0x30 0x15 0x15 0x15 0x15 0x63 0x63 0x63 0x63 0x63 0x63 0x63 0x63
            ##transpose matrix
            a,_,_ = dag.register_block(*permutation_block("TR_#{round_number}", a,trans_p)) # 0x7c, 0x15, 0x63, 0x63, 0x15, 0x63, 0x63, 0x77, 0x63, 0x63, 0xf2, 0x15, 0x63, 0x30, 0x15, 0x63
            #ShiftRows
            a,_,_ = dag.register_block(*shiftrows_block("SR_#{round_number}", a,false))    # 0x7C, 0x15, 0x63, 0x63,
                                                                                           # 0x15, 0x63, 0x63, 0x30,
                                                                                           # 0x63, 0x63, 0xf2, 0x15,
                                                                                           # 0x63, 0x77, 0x15, 0x63
            #MixBytes
            a,_,_ = dag.register_block(*mixcolumns_block("MC_#{round_number}",a, mixcolumns_matrix))
            a,_,_ = dag.register_block(*permutation_block("TR_#{round_number}", a,trans_p))
            ########### key update
            b1,_,_ = dag.register_block(*permutation_block("KP_#{round_number}", b1, tweakey_p))
            b2,_,_ = dag.register_block(*lfsr_block("L_#{round_number}", b2, [
                [1,1,1,1],
                [1,1,1,1],
                [1,1,1,1],
                [1,1,1,1]
            ], tk2_params))
            b2,_,_ = dag.register_block(*permutation_block("KP_#{round_number}", b2, tweakey_p))
            if z == 384
                b3,_,_ = dag.register_block(*lfsr_block("L_#{round_number}", b3, [
                    [1,1,1,1],
                    [1,1,1,1],
                    [1,1,1,1],
                    [1,1,1,1]
                ], tk3_params))
                b3,_,_ = dag.register_block(*permutation_block("KP_#{round_number}", b3, tweakey_p))
            end
        end
        rc,_,_ = dag.register_block(*input_block("RC_#{nb_rounds}", BYTE_RANGE, [4,4], [
            [0x01, 0x02, 0x04, 0x08],    
            [rcon[nb_rounds], rcon[nb_rounds], rcon[nb_rounds], rcon[nb_rounds]],
            [0x00, 0x00, 0x00, 0x00],
            [0x00, 0x00, 0x00, 0x00]
        ]))
        # tweakey xor
        if z == 256
            a,_,_ = dag.register_block(*xor_block("ART_#{nb_rounds}", [a,b1,b2,rc]))
        elsif z == 384
            a,_,_ = dag.register_block(*xor_block("ART_#{nb_rounds}", [a,b1,b2,b3,rc]))
        else
            raise "create_skinny_dag: invalid z parameter (sould be in {256,384}, current value: #{z})"
        end
        dag.set_outputs(a.flatten)
        assert_equal(dag.run(@m+@m+@m), [
            0xc6, 0x92, 0x0e, 0x3c,
            0x82, 0x18, 0xbe, 0x01,
            0x84, 0xbd, 0xc0, 0x1e,
            0x29, 0xd1, 0x9b, 0x01
        ])
    end
end

class TestDeoxys < Minitest::Unit::TestCase
    def test_vectors_256_1()
        dag = create_deoxys_dag(14,256)
        result = dag.run([ # M
            0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00
        ] + [ #TK1
            0x10, 0x11, 0x12, 0x13,   #0x14, 0x15, 0x16, 0x17, #0x28, 0x2a, 0x2c, 0x2e,
            0x14, 0x15, 0x16, 0x17,   #0x19, 0x1a, 0x1b, 0x18, #0x32, 0x34, 0x36, 0x30
            0x18, 0x19, 0x1a, 0x1b,   #0x1e, 0x1f, 0x1c, 0x1d, #0x3c, 0x3e, 0x3
            0x1c, 0x1d, 0x1e, 0x1f    #0x13, 0x10, 0x11, 0x12
        ] + [ #TK2
            0x12, 0x02, 0x12, 0x22,   #0x32, 0x42, 0x52, 0x62, #0x65, 0x84, 0xa4, 0xc5, #
            0x32, 0x42, 0x52, 0x62,   #0x00, 0x00, 0x00, 0x70, #0x00, 0x00, 0x00, 0xe1, #0x19, 0x1a, 0x1b
            0x70, 0x00, 0x00, 0x00,   #0x00, 0x00, 0x00, 0x00, #0x00, 0x00, 0x00, 0x00, #0x1e, 0x1f, 0x1c, 0x1d
            0x00, 0x00, 0x00, 0x00    #0x22, 0x12, 0x02, 0x12, #0x45, 0x24, 0x04, 0x24  #
        ])
        assert_equal(result, [ # C
            0xee, 0xc8, 0x7d, 0xce,
            0x98, 0xd2, 0x9d, 0x40,
            0x78, 0x59, 0x8a, 0xbd,
            0x16, 0xd5, 0x50, 0xff
        ])
    end
   def test_vectors_384_1()
    dag = create_deoxys_dag(16,384)
    result = dag.run([ # M
        0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00
    ] + [ # K1
        0x10, 0x11, 0x12, 0x13,
        0x14, 0x15, 0x16, 0x17,
        0x18, 0x19, 0x1a, 0x1b,
        0x1c, 0x1d, 0x1e, 0x1f
    ] + [ # K2
        0x20, 0x21, 0x22, 0x23,
        0x24, 0x25, 0x26, 0x27,
        0x28, 0x29, 0x2a, 0x2b,
        0x2c, 0x2d, 0x2e, 0x2f
    ] + [ # k3
        0x10, 0x00, 0x10, 0x20,
        0x30, 0x40, 0x50, 0x60,
        0x70, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00
    ])
    assert_equal(result, [ # C
        0x50, 0xb0, 0xde, 0xaa,
        0x3c, 0x31, 0x29, 0xd1,
        0xea, 0x1e, 0xf9, 0x6b,
        0x7c, 0x8d, 0xb6, 0x7f
    ])
end
end