#!/usr/bin/env python3
# Based on https://problemkaputt.de/gbatek.htm#3dscartridgeregisters
import functools
import sys

import ctr_secrets
import cart_crypto

def rotl128(n, r):
    return ((n << r) | (n >> (128 - r))) & ((2**128) - 1)

def derive_key(keyX, keyY):
    x = rotl128(int.from_bytes(keyX, 'big'), 2)
    x = (x ^ int.from_bytes(keyY, 'big')) & ((2**128) - 1)
    x = (x + ctr_secrets.KEY_DERIVATION_CONSTANT) & ((2**128) - 1)
    return rotl128(x, 87).to_bytes(16, 'big')

def card_seed(path):
    with open(path, 'rb') as f:
        cardinfo = f.read(0x1040)

    keyX  = ctr_secrets.BOOTROM_KEYS[0x3b]['keyX']
    keyY  = cardinfo[0x1000:0x1010]
    enc   = cardinfo[0x1010:0x1020]
    tag   = cardinfo[0x1020:0x1030]
    nonce = cardinfo[0x1030:0x103c]

    kek = derive_key(keyX, keyY)
    c = cart_crypto.AES_128_CCM(kek)
    try:
        key = c.decrypt(tag, nonce, enc)
    except ValueError:
        print("Bad authentication tag on title key. Is ctr_secrets correct?")
        sys.exit(1)
    return int.from_bytes(key[0:8], 'big'), int.from_bytes(key[8:16], 'big')

# Use indexes 0-3 depending on ID2
cardrc4keyhalf = ctr_secrets.CART_KEYS[0]['rc4']
cardsnowiv = ctr_secrets.CART_KEYS[0]['snow']

def tolist(val, nbytes):
    return [val >> i & 0xff for i in reversed(range(0, nbytes*8, 8))]

def initsnowstream(snowkey):
    snow20stream = cart_crypto.SNOW20(snowkey, cardsnowiv)
    for _ in range(32): # Discard the first 1024 bits of output
        snow20stream.next()
    return snow20stream

def initrc4stream(snowstream):
    tmp = snowstream.next() << 96 | snowstream.next() << 64 | \
            snowstream.next() << 32 | snowstream.next()
    key = cardrc4keyhalf << 128 | tmp
    rc4stream = cart_crypto.RC4(tolist(key, 32))
    for _ in range(256): # Discard the first 256 bytes of output
        rc4stream.next()
    return rc4stream

# Sample data from https://www.3dbrew.org/wiki/Gamecards
titlekey1, titlekey2 = card_seed('0004000000038c00.3ds')

enccmds = [[0xf32c92d85c9d44ded3e0e41dbe7c90d9, 0x00],
           [0x696b9d8582fb55d31b68cafe70c74a95, 0x04],
           [0xbaa4812ca0ac9c5d19399530e3acccab, 0x04],
           [0x178e427c22d87adb86387249a97d321a, 0x00],
           [0xe06019b1bd5c9130ed6a4d9f4a9e7193, 0x04],
           [0x4e0d224862523bbfe2e6255f80e15f37, 0x04],
           [0x4cdf93d319fb62d0db632a45e3e8d84c, 0x04],
           [0x9aa5d80551002f955546d296a57f0fef, 0x04],
           [0xc12ba81aef30dddbd93fad5d544c6334, 0x04],
           [0x62ec5fb7f420ae1dc6253ae18afa5bb3, 0x200+4], # + 4 byte CRC
           [0xe3fa23aa016be0c93430d1f42ff41324, 0x200+4]] # + 4 byte CRC
# CRC polynomial is 0x82608edb, end result is inverted, CRC over decrypted data

snow20stream = initsnowstream(titlekey1 << 64 | titlekey2)
rc4stream = initrc4stream(snow20stream)

for cmd, datalen in enccmds:
    print('encrypted command: %032x' % cmd)
    dec = functools.reduce(lambda x, y: x<<8 | y, [
        rc4stream.next() ^ x for x in tolist(cmd, 16)])
    print('decrypted command: %032x' % dec)

    if dec >> (15*8) == 0x83: #"set seed" command
        randseed = dec & 0xFFFFFFFFFFFFFFFF
        print('changing "seed" to %016x' % randseed)
        snow20stream = initsnowstream(titlekey1 << 64 | randseed)

    rc4stream = initrc4stream(snow20stream)

    print('skipping %d bytes' % datalen)
    for _ in range(datalen):
        rc4stream.next() #these RC4 bytes would be used to decrypt the data