go-pkparse-server/pokemon-parsing/gen7/parse.go

package gen7

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"strings"
	"unicode/utf16"

	p "github.com/ajswis/go-pkparse-server/pokemon"
)

// Parse accepts a raw pokemon (byte slice) and either returns a Pokemon struct
// with all applicable fields populated or an error depicting the reason for
// failure.
func Parse(rawPokemon p.RawPokemon) (*p.Pokemon, error) {
	if err := validateRawPokemon(rawPokemon); err != nil {
		return nil, err
	}

	var pkmn p.Pokemon
	pkmn.RawPokemon = &rawPokemon

	pkmn.EncryptionConstant = binary.LittleEndian.Uint32(rawPokemon[0x00:0x08])
	pkmn.PokedexNumber = toUint16(rawPokemon[0x08:0x0a])
	pkmn.HeldItemID = toUint16(rawPokemon[0x0a:0x0c])
	pkmn.TrainerID = toUint16(rawPokemon[0x0c:0x0e])
	pkmn.SecretID = toUint16(rawPokemon[0x0e:0x10])
	pkmn.FullTrainerID = binary.LittleEndian.Uint32(rawPokemon[0x0c:0x10])
	pkmn.Experience = binary.LittleEndian.Uint32(rawPokemon[0x10:0x15])
	pkmn.AbilityID = rawPokemon[0x14]
	pkmn.AbilityNum = rawPokemon[0x14]
	pkmn.PID = binary.LittleEndian.Uint32(rawPokemon[0x18:0x1c])
	pkmn.NatureID = rawPokemon[0x1c]

	genderByte := rawPokemon[0x1d]
	pkmn.FatefulEncounter = (genderByte & 0x01) == 1
	if (genderByte & 0x02) == 0x02 {
		pkmn.Gender = "F"
	} else if (genderByte & 0x04) == 0x04 {
		pkmn.Gender = ""
	} else {
		pkmn.Gender = "M"
	}
	pkmn.FormID = genderByte >> 3

	pkmn.HPEV = rawPokemon[0x1e]
	pkmn.AtkEV = rawPokemon[0x1f]
	pkmn.DefEV = rawPokemon[0x20]
	pkmn.SpeEV = rawPokemon[0x21]
	pkmn.SpAtkEV = rawPokemon[0x22]
	pkmn.SpDefEV = rawPokemon[0x23]

	pokerusByte := rawPokemon[0x2b]
	pkmn.PokerusDuration = pokerusByte & 15
	pkmn.PokerusStrain = pokerusByte >> 4

	pkmn.RawNickname = rawPokemon[0x40:0x58]
	pkmn.Nickname = binaryToUTF16leString(pkmn.RawNickname)

	pkmn.Move1ID = toUint16(rawPokemon[0x5a:0x5c])
	pkmn.Move2ID = toUint16(rawPokemon[0x5c:0x5e])
	pkmn.Move3ID = toUint16(rawPokemon[0x5e:0x60])
	pkmn.Move4ID = toUint16(rawPokemon[0x60:0x62])
	pkmn.Move1PP = rawPokemon[0x62]
	pkmn.Move2PP = rawPokemon[0x63]
	pkmn.Move3PP = rawPokemon[0x64]
	pkmn.Move4PP = rawPokemon[0x65]
	pkmn.Move1PPUsed = rawPokemon[0x66]
	pkmn.Move2PPUsed = rawPokemon[0x67]
	pkmn.Move3PPUsed = rawPokemon[0x68]
	pkmn.Move4PPUsed = rawPokemon[0x69]
	pkmn.EggMove1ID = toUint16(rawPokemon[0x6a:0x6c])
	pkmn.EggMove2ID = toUint16(rawPokemon[0x6c:0x6e])
	pkmn.EggMove3ID = toUint16(rawPokemon[0x6e:0x70])
	pkmn.EggMove4ID = toUint16(rawPokemon[0x70:0x72])

	ivBytes := binary.LittleEndian.Uint32(rawPokemon[0x74:0x78])
	pkmn.HPIV = ivBytes & 0x1f
	pkmn.AtkIV = ivBytes >> 5 & 0x1f
	pkmn.DefIV = ivBytes >> 10 & 0x1f
	pkmn.SpeIV = ivBytes >> 15 & 0x1f
	pkmn.SpAtkIV = ivBytes >> 20 & 0x1f
	pkmn.SpDefIV = ivBytes >> 25 & 0x1f
	pkmn.IsEgg = ((ivBytes >> 30) % 2) == 1
	pkmn.IsNicknamed = ((ivBytes >> 31) % 2) == 1

	return &pkmn, nil
}

func binaryToUTF16leString(b []byte) string {
	r := bytes.NewReader(b)
	u16 := make([]uint16, len(b)/2)
	binary.Read(r, binary.LittleEndian, &u16)
	runes := utf16.Decode(u16)
	s := strings.Split(string(runes), "\u0000")[0]

	return s
}

func calculateChecksum(b []byte) uint16 {
	var sum uint16
	for i := 8; i < 232; i += 2 {
		sum += toUint16(b[i : i+2])
	}
	return sum & 0xffff
}

func toUint16(b []byte) uint16 {
	return binary.LittleEndian.Uint16(b)
}

func validateRawPokemon(rawPokemon p.RawPokemon) error {
	if l := len(rawPokemon); l != 232 && l != 260 {
		return fmt.Errorf("Invalid length for generation 7 pokemon: %d bytes", l)
	}

	checksumValue := toUint16(rawPokemon[0x06:0x08])
	if calculatedChecksum := calculateChecksum(rawPokemon); calculatedChecksum != checksumValue {
		return fmt.Errorf("Invalid checksum for generation 7 pokemon")
	}

	if toUint16(rawPokemon[0x04:0x06]) == 1 || rawPokemon[0x58] == 1 || rawPokemon[0x90] == 1 || rawPokemon[0xc8] == 1 {
		return fmt.Errorf("Invalid bytes supplied for generation 7 pokemon")
	}

	return nil
}