# aes-xts-pur64 is OpenCL code for aes-xts256-plain64 encryption compatible with LUKS
#
#     Copyright (C) 2025  Kirill Shakirov
#
#     This program is free software: you can redistribute it and/or modify
#     it under the terms of the GNU General Public License as published by
#     the Free Software Foundation, either version 3 of the License, or
#     (at your option) any later version.
#
#     This program is distributed in the hope that it will be useful,
#     but WITHOUT ANY WARRANTY; without even the implied warranty of
#     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#     GNU General Public License for more details.
#
#     You should have received a copy of the GNU General Public License
#     along with this program.  If not, see <https://www.gnu.org/licenses/>.

import json
import subprocess
import nyanger.simple.static as nya_stat
import random

# Init logger
log = nya_stat.get_logger("nyan")

LUKS_VOL_FILE_NAME = "./data/test_vol.img"
UNENCRYPT_DATA_FILE_NAME= "./data/unencrypt.img"
OUT_FILE_NAME = "./data/test_vectors.json"
KEY_FILE_NAME = "./data/master.key"
VECTORS_NUM=256

def read_metadata(file_name: str) -> dict:
    #cryptsetup luksDump --dump-json-metadata /dev/loop0
    luks_cmd: list[str] = ["cryptsetup", "luksDump", "--dump-json-metadata", file_name]

    result = subprocess.run(luks_cmd, capture_output=True, encoding="UTF-8")
    if result.returncode == 0 and result.stdout is not None:
        metadata = json.loads(result.stdout)
        return metadata
    else:
        raise Exception(f"Error executing 'cryptsetup' binary! {result.stderr}")


def main():
    log.info("Start :3")
    metad = read_metadata(LUKS_VOL_FILE_NAME)
    log.info(f"metadata:\n{metad}")

    segments_offset_bytes = int(metad["segments"]["0"]["offset"])

    with open(LUKS_VOL_FILE_NAME, 'rb') as luks_file:
        luks_file.seek(segments_offset_bytes)
        enc_data = luks_file.read(512*VECTORS_NUM)

    with open(UNENCRYPT_DATA_FILE_NAME, 'rb') as data_file:
        uenc_data = data_file.read(512*VECTORS_NUM)

    with open(KEY_FILE_NAME, "rb") as key_file:
        key_data = key_file.read(32)

    data_list = []
    for sec_n in range(VECTORS_NUM):
        block_num = random.randint(0, 31)
        _offset = ((sec_n*512)+(block_num*16))
        entry = {
            "sector_number": sec_n,
            "block_number": block_num,
            "unencrypted_data": uenc_data[_offset:_offset+16].hex(),
            "encrypted_data": enc_data[_offset:_offset+16].hex()
        }
        data_list.append(entry)

    with open(OUT_FILE_NAME, 'wt') as out_file:
        json.dump({"encryption_key": key_data.hex(sep=' ', bytes_per_sep=16), "vectors": data_list},
                  out_file, indent="  ")

    log.info(" --- Finish --- ")


if __name__ == '__main__':
    main()