#This script will generate dat files that allow offline play of premium modules for the game Neverwinter Nights.
#YOU MUST ONLY USE THIS SCRIPT TO GENERATE DATS FOR PREMIUM MODULES THAT YOU HAVE PURCHASED!
import os, sys, socket, struct
from binascii import unhexlify

try:
    from Crypto.Cipher import Blowfish
except ImportError:
    print 'You must have PyCrypto installed to use this script.'
    print "\nIf you're on Windows, the easiest way to get it is to grab an installer from:"
    print 'http://www.voidspace.org.uk/python/modules.shtml#pycrypto'
    sys.exit()


#These are only part of the key used to decrypt the module.
#The game generates the last 16 bytes of the key from the MD5 hash over the NWM file for the module.
#Once you have the full key, you can decrypt the HAK with Blowfish in ECB mode, ignoring the first 8 bytes of the file.

#Edit: Removed the keys, the script just grabs them from the master server now.
#But there should be enough info in this file to work out how to decrypt the dat and get them.
modules = [
    b"Neverwinter Nights - Kingmaker.dat",
    b"Neverwinter Nights - ShadowGuard.dat",
    b"Neverwinter Nights - Witch's Wake.dat",
    b"Neverwinter Nights - Pirates of the Sword Coast.dat",
    b"Neverwinter Nights - Wyvern Crown of Cormyr.dat",
    b"Neverwinter Nights - Infinite Dungeons.dat"
    ]

if len(sys.argv) < 2:
    print 'Usage: %s *Key1*' % (os.path.basename(sys.argv[0]))
    print 'Example: %s ABCDE-FGHIJ-KLMNO-PQRST-UVWXY-ZABCD-EFGHI' % (os.path.basename(sys.argv[0]))
    print '\nTo find your CD key, look in "nwncdkey.ini" in the game folder.'
    print 'You want the key listed as "Key1".'
    print'\nYOU MUST ONLY USE THIS SCRIPT TO GENERATE DATS\nFOR PREMIUM MODULES THAT YOU HAVE PURCHASED!'
    sys.exit()

cdKey = sys.argv[1].upper()
cdKey = cdKey.replace('-','')
if len(cdKey) != 35:
    print 'The CD key(minus dashes) should be 35 characters long'
    sys.exit()
cdKey = cdKey + b'\x00' * (56-len(cdKey)) #Pad the CD key on the right with 0x00 to 56 bytes to get the encryption key for the dat.

cipher = Blowfish.new(cdKey, Blowfish.MODE_ECB)

print '\nModules:'
for idx, module in enumerate(modules):
    print '[%u] %s' % (idx+1, module[:-4])
print '[q] quit'
print'\nYOU MUST ONLY USE THIS SCRIPT TO GENERATE DATS\nFOR PREMIUM MODULES THAT YOU HAVE PURCHASED!\n'

sock1 = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock1.bind(('', 5120))

while(1):
    inp = raw_input("Module choice: ")
    try:
        mnum = int(inp)-1
    except ValueError:
        if inp == 'q':
            sys.exit()
        continue
    if (mnum < 0) or (mnum > 5):
        continue

    #The master server currently lets any properly formatted auth request through, regardless of data.
    reqpckt = 'BMDQ'
    reqpckt += struct.pack('<H', 0x1400) #Port we're listening on
    reqpckt += struct.pack('<H', len(modules[mnum][:-4])) #Length of module name
    reqpckt += modules[mnum][:-4] #Module name
    reqpckt += struct.pack('<H', 0x20)
    reqpckt += b'\x00' * 0x20 #No clue what this is, but it's unchecked by the server currently.
    reqpckt += struct.pack('<H', 0x01) #Number of times the next block will happen(the block consists of the 0x8 byte + 0x20 byte stuff)
    reqpckt += struct.pack('<H', 0x08)
    reqpckt += 'DEADBEA7' #Public CD key, unchecked by the server currently.
    reqpckt += struct.pack('<H', 0x20)
    reqpckt += b'\xFF' * 0x20 #Supposed to be some hash or something, unchecked by the server currently(but can't be 0x00..00).
    reqpckt += struct.pack('<H', 0x44) #Size of public RSA key
    reqpckt += struct.pack('<I', 0x01) #Exponent
    reqpckt += b'\xFF' * 0x40 #Modulus
    #Basically, they'll compute Msg^1 mod 0xFF...FF, which will just give us the original unencrypted message.
    #That way, we don't have to deal with any RSA crap. :)

    print 'Sending auth request...'
    sock1.sendto(reqpckt, ("nwmaster.bioware.com", 5121))
    response, addr = sock1.recvfrom(2048)

    if response[:5] != b'BMDR\x01':
        print 'Unexpected server response!'
        continue

    datbuff = struct.pack('<B', 0x01) #Pointer to key data(just this 1 byte)
    datbuff += response[0x21:0x21+0x28] #Key data
    datbuff += b'\x00' * (1024-1-len(datbuff)) #Rest of the file can be whatever, it's not checked by the game.
    datbuff += struct.pack('<B', 0x28) #Hardcoded check in the game for 0x28 here

    if len(datbuff) != 1024:
        print 'Error'
        sys.exit()

    with open(modules[mnum], 'wb') as fh:
        fh.write(cipher.encrypt(datbuff)) #Encrypt the dat data and write it out.

    print 'Created "%s"!' % (modules[mnum])