ketti

/**
 * Java implementation of the secret password check and the You Win code generation algorithm
 * in the NES game "Treasure Master"
 *
 * Note: The password check is possibly incomplete.
 *
 * @author ketti
 * @see http://www.reddit.com/r/TreasureMaster/
 */
public class TreasureMaster
{
    private static final int YOUWINCODE_LENGTH = 0x18;
    private static final int PASSWORD_LENGTH = 0x18;
    private static final int BUFFER_LENGTH = 0x0F;
    private static final int CHECK_LENGTH_1 = 0x0D;
    private static final int CHECK_LENGTH_2 = 0x08;
    private static final int SERIAL_LENGTH = 0x08;

    private static final byte[] PPU = new byte[] {(byte)0xFD, (byte)0x22, (byte)0x3C, (byte)0x40};

    private static char[] CONVERSION_TABLE = new char[] {
            '0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
            'B', 'C', 'D', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N',
            'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'X', 'Y', 'Z', '!'};

    public static void main(String[] args)
    {
        String secret = "3HDJL9DNQV2WYTV4S91RXR86";
        System.out.println("Secret password \"" + secret + "\" is " +
                (checkSecretPassword(secret) ? "" : "in") + "valid.\n");



        String serial = "12345678";
        int score = 79900;
        int lives = 4;

        System.out.println("Serial: " + serial);
        System.out.println("Score: " + score);
        System.out.println("Lives: " + lives);

        System.out.println("\nYou Win codes:");
        for (int random = 0; random < 0x10; random++)
        {
            System.out.println(computeYouWinCode(serial, score, lives, random));
        }
    }


    public static boolean checkSecretPassword(String string)
    {
        if (string.length() != PASSWORD_LENGTH)
        {
            throw new IllegalArgumentException("Secret password must be " +
                    PASSWORD_LENGTH + " characters long.");
        }

        // Convert the secret password string to the "Treasure Master format".
        byte[] password = new byte[PASSWORD_LENGTH];
        char[] chars = string.toCharArray();
        for (int i = 0; i < PASSWORD_LENGTH; i++)
        {
            boolean found = false;
            for (int j = 0; j < CONVERSION_TABLE.length; j++)
            {
                if (chars[i] == CONVERSION_TABLE[j])
                {
                    password[i] = (byte)j;
                    found = true;
                    break;
                }
            }
            if (!found)
            {
                throw new IllegalArgumentException("Illegal character in secret password found: " + chars[i]);
            }
        }


        // Compute buffer contents for the first check (somewhere around $B0BE)
        byte[] buffer = new byte[BUFFER_LENGTH];
        for (int i = 0; i < PASSWORD_LENGTH; i++)
        {
            byte b = password[i];

            CarryContainer c = new CarryContainer();
            b = c.shiftLeftWithCarry(b);
            b = c.shiftLeftWithCarry(b);
            b = c.shiftLeftWithCarry(b);

            for (int k = 0; k < 5; k++)
            {
                b = c.rotateLeftThroughCarry(b);

                for (int j = BUFFER_LENGTH - 1; j >= 0; j--)
                {
                    buffer[j] = c.rotateLeftThroughCarry(buffer[j]);
                }
            }
        }

        int sum = computeChecksum(buffer, CHECK_LENGTH_1);

        byte lowByte = (byte)(sum & 0xFF);
        byte highByte = (byte)((sum >> 8) & 0xFF);

        // Compare computed checksum with checksum bytes in the buffer
        if ((lowByte != buffer[CHECK_LENGTH_1]) || (highByte != buffer[CHECK_LENGTH_1 + 1]))
        {
            return false;
        }

        // For the second check to execute the first byte of the buffer has to be smaller than 0x20
        if (buffer[0] >= 0x20)
        {
            return false;
        }

        // Compute buffer contents for the second check (somewhere around $B0EB)
        CarryContainer c = new CarryContainer();
        for (int i = buffer[0]; i > 0; i--)
        {
            for (int j = buffer.length - 1; j >= 0; j--)
            {
                buffer[j] = c.rotateLeftThroughCarry(buffer[j]);
            }
        }

        sum = computeChecksum(buffer, CHECK_LENGTH_2);

        return ((byte)(sum & 0xFF)) == buffer[8];
    }

    public static String computeYouWinCode(final String serial, final int score,
            final int lives, final int random)
    {
        byte[] serialBytes = new byte[SERIAL_LENGTH];
        byte[] scoreBytes = new byte[SERIAL_LENGTH];
        byte[] buffer = new byte[BUFFER_LENGTH];

        if (serial.length() != SERIAL_LENGTH)
        {
            throw new IllegalArgumentException("Serial must be " + SERIAL_LENGTH + " characters long.");
        }

        // Convert serial to byte array
        char[] serialChars = serial.toCharArray();
        for (int i = 0; i < SERIAL_LENGTH; i++)
        {
            char c = serialChars[i];
            if ((c < '0') || (c > '9'))
            {
                throw new IllegalArgumentException("Invalid character in serial: " + c);
            }

            serialBytes[i] = (byte)(c - '0');
        }

        // Copy PPU bytes to the buffer
        for (int i = 0; i < PPU.length; i++)
        {
            buffer[i] = PPU[i];
        }

        // Convert score to "Treasure Master format"
        int s = score / 10;
        for (int i = SERIAL_LENGTH - 1; i >= 0 ; i--)
        {
            scoreBytes[i] = (byte)(s % 10);
            s /= 10;
        }

        // Put score and serial into the buffer
        byte b;
        CarryContainer c = new CarryContainer();
        for (int i = 0; i < SERIAL_LENGTH; i++)
        {
            b = scoreBytes[i];
            b = c.shiftLeftWithCarry(b);
            b = c.shiftLeftWithCarry(b);
            b = c.shiftLeftWithCarry(b);
            b = c.shiftLeftWithCarry(b);
            buffer[i + 4] = (byte)(serialBytes[i] | b);
        }

        // Put lives and random number into the buffer
        b = (byte)lives;
        b = c.shiftLeftWithCarry(b);
        b = c.shiftLeftWithCarry(b);
        b = c.shiftLeftWithCarry(b);
        b = c.shiftLeftWithCarry(b);
        buffer[4 + SERIAL_LENGTH] = (byte)((random & 0x0F) | b);

        // Compute checksum for that buffer and append it to the buffer
        int sum = computeChecksum(buffer, CHECK_LENGTH_1);
        buffer[CHECK_LENGTH_1] = (byte)(sum & 0xFF);
        buffer[CHECK_LENGTH_1 + 1] = (byte)((sum >> 8) & 0xFF);

        // Create You Win code string
        StringBuffer sb = new StringBuffer(YOUWINCODE_LENGTH);
        c.setCarry(false);
        for (int i = 0; i < YOUWINCODE_LENGTH; i++)
        {
            b = 0;

            for (int j = 0; j < 5; j++)
            {
                for (int k = BUFFER_LENGTH - 1; k >= 0; k--)
                {
                    buffer[k] = c.rotateLeftThroughCarry(buffer[k]);
                }
                b = c.rotateLeftThroughCarry(b);
            }

            sb.insert(i, CONVERSION_TABLE[b]);
        }

        return sb.toString();
    }

    private static int computeChecksum(final byte[] buffer, final int length)
    {
        int sum = 0;
        for (int i = 0; i < length; i++)
        {
            sum += buffer[i] & 0xFF;
        }

        return sum;
    }

    static class CarryContainer
    {
        private boolean carry;

        public CarryContainer()
        {
            this(false);
        }

        public CarryContainer(final boolean carry)
        {
            this.carry = carry;
        }

        public byte shiftLeftWithCarry(final byte b)
        {
            // ASL
            carry = ((b & 0x80) != 0);
            return (byte)(b << 1);
        }

        public byte rotateLeftThroughCarry(byte b)
        {
            // ROL
            boolean newCarry = ((b & 0x80) != 0);
            b = (byte)((b << 1) | ((carry) ? 1 : 0));
            carry = newCarry;

            return b;
        }

        public void setCarry(boolean carry)
        {
            this.carry = carry;
        }
    }
}