from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad


def save_bgroup(data):
    with open("bgroupcert.dat", "wb") as f:
        f.write(data)


def save_zgpriv(data):
    with open("zgpriv.dat", "wb") as f:
       f.write(data[:32])


def save_keybox(data):
    with open("keybox.dat", "wb") as f:
       f.write(data[:128])



def decrypt_aes_ecb(input_file, key):
    with open(input_file, "rb") as f:
        encrypted_data = f.read()
    cipher = AES.new(bytes.fromhex(key), AES.MODE_ECB)
    dec_data = cipher.decrypt(pad(encrypted_data, 16))
    if b"INNER_MSTAR" in dec_data:
        fixed_data = dec_data[64:]
        if b"CHAI" in fixed_data:
            print("BGROUPCERT.dat SAVED AS bgroupcert.dat")
            save_bgroup(fixed_data)
        elif b"kbox" in fixed_data:
            print("WIDEVINE KEYBOX SAVED AS keybox.dat")
            save_keybox(fixed_data)
        else:
            print("ZGPRIV saved as zgpriv.dat")
            save_zgpriv(fixed_data)
    else:
        print("Key not working")

key = "0007FF4154534D92FC55AA0FFF0110E0"
decrypt_aes_ecb("keybox_encrypted.dat", key)