// online tester: https://lzltool.cn/SM4
import x.crypto.sm4

struct SM4Data {
    key    string
    iv     string
    input  string
    output string
}

const sm4data_ecb = [
    SM4Data{
        key:    '0123456789abcdeffedcba9876543210'
        input:  '0123456789abcdeffedcba9876543210'
        output: '681edf34d206965e86b3e94f536e4246'
    },
    SM4Data{
        key:    '0123456789abcdeffedcba9876543210'
        input:  '00112233445566778899aabbccddeeff'
        output: '09325c4853832dcb9337a5984f671b9a'
    },
    SM4Data{
        key:    '456789abcdeffedcba98765432100123'
        input:  '2233445566778899aabbccddeeff0011'
        output: '58ab414d84fb3008b0bee987f97021e6'
    },
    SM4Data{
        key:    '89abcdeffedcba987654321001234567'
        input:  '445566778899aabbccddeeff00112233445566778899aabbccddeeff00112233'
        output: '5937a929a2d9137216c72a28cd9cf6195937a929a2d9137216c72a28cd9cf619'
    },
]

const sm4data_cbc = [
    SM4Data{
        key:    '0123456789abcdeffedcba9876543210'
        iv:     '0123456789abcdeffedcba9876543210'
        input:  '0123456789abcdeffedcba9876543210'
        output: '2677f46b09c122cc975533105bd4a22a'
    },
    SM4Data{
        key:    '0123456789abcdeffedcba9876543210'
        iv:     '2677f46b09c122cc975533105bd4a22a'
        input:  '00112233445566778899aabbccddeeff'
        output: '2e7b41173b42b90189bbcf4eeb5b5145'
    },
    SM4Data{
        key:    '0123456789abcdeffedcba9876543210'
        iv:     '2e7b41173b42b90189bbcf4eeb5b5145'
        input:  '29860013b6b74c7503c524b62ecd52df'
        output: 'cd0cd15a934b3f47476292113d7f35bd'
    },
]

fn test_sm4_ecb() {
    mut key := []u8{}
    mut input := []u8{}
    mut output := []u8{}

    mut cipher := &sm4.SM4Cipher{}

    for data in sm4data_ecb {
        key = '0x${data.key}'.u8_array()
        input = '0x${data.input}'.u8_array()
        output = [u8(0)].repeat(input.len) // output must be exactly the same length as input
        assert key.len == 16
        cipher = sm4.new_cipher(.sm4_encrypt, key)!
        cipher.crypt_ecb(input, mut output)!
        assert output.hex() == data.output

        cipher = sm4.new_cipher(.sm4_decrypt, key)!
        cipher.crypt_ecb(output, mut output)!
        assert output.hex() == data.input
    }
}

fn test_sm4_cbc() {
    mut key := []u8{}
    mut iv := []u8{}
    mut input := []u8{}
    mut output := []u8{}
    mut iv_next := []u8{}

    mut cipher := &sm4.SM4Cipher{}

    // CBC encrypt
    for i, data in sm4data_cbc {
        key = '0x${data.key}'.u8_array()
        iv = '0x${data.iv}'.u8_array()
        input = '0x${data.input}'.u8_array()
        output = [u8(0)].repeat(input.len) // output must be exactly the same length as input
        assert key.len == 16
        assert iv.len == 16
        if i != 0 {
            assert iv_next == iv
        }
        cipher = sm4.new_cipher(.sm4_encrypt, key)!
        iv_next = cipher.crypt_cbc(iv, input, mut output)!
        dump(iv_next.hex())
        assert output.hex() == data.output
    }
    // CBC decrypt
    for i, data in sm4data_cbc {
        key = '0x${data.key}'.u8_array()
        iv = '0x${data.iv}'.u8_array()
        input = '0x${data.output}'.u8_array()
        output = [u8(0)].repeat(input.len) // output must be exactly the same length as input
        assert key.len == 16
        assert iv.len == 16
        if i != 0 {
            assert iv_next == iv
        }
        cipher = sm4.new_cipher(.sm4_decrypt, key)!
        iv_next = cipher.crypt_cbc(iv, input, mut output)!
        dump(iv_next.hex())
        assert output.hex() == data.input
    }
}

fn test_sm4_wrong_length() ! {
    // wrong key length test
    mut fail_flag := false
    sm4.new_cipher(.sm4_encrypt, [u8(0xff)].repeat(33)) or {
        fail_flag = true
        assert err.msg() == 'SM4 only support 128bit key length'
    }
    assert fail_flag

    fail_flag = false
    sm4.new_cipher(.sm4_decrypt, [u8(0xff)].repeat(33)) or {
        fail_flag = true
        assert err.msg() == 'SM4 only support 128bit key length'
    }
    assert fail_flag

    // wrong input length test(ecb)
    mut output := []u8{len: 32}
    fail_flag = false
    mut c1 := sm4.new_cipher(.sm4_encrypt, [u8(0xff)].repeat(16))!
    c1.crypt_ecb([u8(0xff)].repeat(111), mut output) or {
        fail_flag = true
        assert err.msg() == 'input must be padded to multiple of 16 bytes'
    }
    assert fail_flag

    // wrong output length test(ecb)
    fail_flag = false
    mut c2 := sm4.new_cipher(.sm4_encrypt, [u8(0xff)].repeat(16))!
    c2.crypt_ecb([u8(0xff)].repeat(16), mut output) or {
        fail_flag = true
        assert err.msg() == 'output must be exactly the same length as input'
    }
    assert fail_flag

    // wrong iv length test(cbc)
    fail_flag = false
    mut c3 := sm4.new_cipher(.sm4_encrypt, [u8(0xff)].repeat(16))!
    c3.crypt_cbc([u8(0xff)].repeat(22), [u8(0xff)].repeat(16), mut [u8(0xff)].repeat(16)) or {
        fail_flag = true
        assert err.msg() == 'iv length must be exactly 16 bytes'
    }
    assert fail_flag

    // wrong input length test(cbc)
    fail_flag = false
    mut c4 := sm4.new_cipher(.sm4_encrypt, [u8(0xff)].repeat(16))!
    c4.crypt_cbc([u8(0xff)].repeat(16), [u8(0xff)].repeat(111), mut output) or {
        fail_flag = true
        assert err.msg() == 'input must be padded to multiple of 16 bytes'
    }
    assert fail_flag

    // wrong output length test(cbc)
    fail_flag = false
    mut c5 := sm4.new_cipher(.sm4_encrypt, [u8(0xff)].repeat(16))!
    c5.crypt_cbc([u8(0xff)].repeat(16), [u8(0xff)].repeat(16), mut output) or {
        fail_flag = true
        assert err.msg() == 'output must be exactly the same length as input'
    }
    assert fail_flag
}