256 lines
7.4 KiB
Go
256 lines
7.4 KiB
Go
package MPQ
|
|
|
|
import (
|
|
"bytes"
|
|
"compress/zlib"
|
|
"encoding/binary"
|
|
"fmt"
|
|
"log"
|
|
"strings"
|
|
|
|
"github.com/JoshVarga/blast"
|
|
"github.com/OpenDiablo2/OpenDiablo2/Common"
|
|
"github.com/OpenDiablo2/OpenDiablo2/Compression"
|
|
)
|
|
|
|
// Stream represents a stream of data in an MPQ archive
|
|
type Stream struct {
|
|
MPQData MPQ
|
|
BlockTableEntry BlockTableEntry
|
|
FileName string
|
|
EncryptionSeed uint32
|
|
BlockPositions []uint32
|
|
CurrentPosition uint32
|
|
CurrentData []byte
|
|
CurrentBlockIndex uint32
|
|
BlockSize uint32
|
|
}
|
|
|
|
// CreateStream creates an MPQ stream
|
|
func CreateStream(mpq MPQ, blockTableEntry BlockTableEntry, fileName string) *Stream {
|
|
result := &Stream{
|
|
MPQData: mpq,
|
|
BlockTableEntry: blockTableEntry,
|
|
CurrentBlockIndex: 0xFFFFFFFF,
|
|
}
|
|
fileSegs := strings.Split(fileName, `\`)
|
|
result.EncryptionSeed = hashString(fileSegs[len(fileSegs)-1], 3)
|
|
if result.BlockTableEntry.HasFlag(FileFixKey) {
|
|
result.EncryptionSeed = (result.EncryptionSeed + result.BlockTableEntry.FilePosition) ^ result.BlockTableEntry.UncompressedFileSize
|
|
}
|
|
result.BlockSize = 0x200 << result.MPQData.Data.BlockSize
|
|
|
|
if result.BlockTableEntry.HasFlag(FilePatchFile) {
|
|
log.Fatal("Patching is not supported")
|
|
}
|
|
|
|
if (result.BlockTableEntry.HasFlag(FileCompress) || result.BlockTableEntry.HasFlag(FileImplode)) && !result.BlockTableEntry.HasFlag(FileSingleUnit) {
|
|
result.loadBlockOffsets()
|
|
}
|
|
return result
|
|
}
|
|
|
|
func (v *Stream) loadBlockOffsets() {
|
|
blockPositionCount := ((v.BlockTableEntry.UncompressedFileSize + v.BlockSize - 1) / v.BlockSize) + 1
|
|
v.BlockPositions = make([]uint32, blockPositionCount)
|
|
v.MPQData.File.Seek(int64(v.BlockTableEntry.FilePosition), 0)
|
|
binary.Read(v.MPQData.File, binary.LittleEndian, &v.BlockPositions)
|
|
blockPosSize := blockPositionCount << 2
|
|
if v.BlockTableEntry.HasFlag(FileEncrypted) {
|
|
decrypt(v.BlockPositions, v.EncryptionSeed-1)
|
|
if v.BlockPositions[0] != blockPosSize {
|
|
panic("Decryption of MPQ failed!")
|
|
}
|
|
if v.BlockPositions[1] > v.BlockSize+blockPosSize {
|
|
panic("Decryption of MPQ failed!")
|
|
}
|
|
}
|
|
}
|
|
|
|
func (v *Stream) Read(buffer []byte, offset, count uint32) uint32 {
|
|
if v.BlockTableEntry.HasFlag(FileSingleUnit) {
|
|
return v.readInternalSingleUnit(buffer, offset, count)
|
|
}
|
|
toRead := count
|
|
readTotal := uint32(0)
|
|
for toRead > 0 {
|
|
read := v.readInternal(buffer, offset, count)
|
|
if read == 0 {
|
|
break
|
|
}
|
|
readTotal += read
|
|
offset += read
|
|
toRead -= read
|
|
}
|
|
return readTotal
|
|
}
|
|
|
|
func (v *Stream) readInternalSingleUnit(buffer []byte, offset, count uint32) uint32 {
|
|
if len(v.CurrentData) == 0 {
|
|
v.loadSingleUnit()
|
|
}
|
|
|
|
bytesToCopy := Common.Min(uint32(len(v.CurrentData))-v.CurrentPosition, count)
|
|
copy(buffer[offset:offset+bytesToCopy], v.CurrentData[v.CurrentPosition:v.CurrentPosition+bytesToCopy])
|
|
v.CurrentPosition += bytesToCopy
|
|
return bytesToCopy
|
|
}
|
|
|
|
func (v *Stream) readInternal(buffer []byte, offset, count uint32) uint32 {
|
|
v.bufferData()
|
|
localPosition := v.CurrentPosition % v.BlockSize
|
|
bytesToCopy := Common.MinInt32(int32(len(v.CurrentData))-int32(localPosition), int32(count))
|
|
if bytesToCopy <= 0 {
|
|
return 0
|
|
}
|
|
copy(buffer[offset:offset+uint32(bytesToCopy)], v.CurrentData[localPosition:localPosition+uint32(bytesToCopy)])
|
|
v.CurrentPosition += uint32(bytesToCopy)
|
|
return uint32(bytesToCopy)
|
|
}
|
|
|
|
func (v *Stream) bufferData() {
|
|
requiredBlock := uint32(v.CurrentPosition / v.BlockSize)
|
|
if requiredBlock == v.CurrentBlockIndex {
|
|
return
|
|
}
|
|
expectedLength := Common.Min(v.BlockTableEntry.UncompressedFileSize-(requiredBlock*v.BlockSize), v.BlockSize)
|
|
v.CurrentData = v.loadBlock(requiredBlock, expectedLength)
|
|
v.CurrentBlockIndex = requiredBlock
|
|
}
|
|
|
|
func (v *Stream) loadSingleUnit() {
|
|
fileData := make([]byte, v.BlockSize)
|
|
v.MPQData.File.Seek(int64(v.MPQData.Data.HeaderSize), 0)
|
|
binary.Read(v.MPQData.File, binary.LittleEndian, &fileData)
|
|
if v.BlockSize == v.BlockTableEntry.UncompressedFileSize {
|
|
v.CurrentData = fileData
|
|
return
|
|
}
|
|
v.CurrentData = decompressMulti(fileData, v.BlockTableEntry.UncompressedFileSize)
|
|
}
|
|
|
|
func (v *Stream) loadBlock(blockIndex, expectedLength uint32) []byte {
|
|
var (
|
|
offset uint32
|
|
toRead uint32
|
|
)
|
|
if v.BlockTableEntry.HasFlag(FileCompress) || v.BlockTableEntry.HasFlag(FileImplode) {
|
|
offset = v.BlockPositions[blockIndex]
|
|
toRead = v.BlockPositions[blockIndex+1] - offset
|
|
} else {
|
|
offset = blockIndex * v.BlockSize
|
|
toRead = expectedLength
|
|
}
|
|
offset += v.BlockTableEntry.FilePosition
|
|
data := make([]byte, toRead)
|
|
v.MPQData.File.Seek(int64(offset), 0)
|
|
binary.Read(v.MPQData.File, binary.LittleEndian, &data)
|
|
if v.BlockTableEntry.HasFlag(FileEncrypted) && v.BlockTableEntry.UncompressedFileSize > 3 {
|
|
if v.EncryptionSeed == 0 {
|
|
panic("Unable to determine encryption key")
|
|
}
|
|
|
|
decryptBytes(data, blockIndex+v.EncryptionSeed)
|
|
}
|
|
if v.BlockTableEntry.HasFlag(FileCompress) && (toRead != expectedLength) {
|
|
if !v.BlockTableEntry.HasFlag(FileSingleUnit) {
|
|
data = decompressMulti(data, expectedLength)
|
|
} else {
|
|
data = pkDecompress(data)
|
|
}
|
|
}
|
|
if v.BlockTableEntry.HasFlag(FileImplode) && (toRead != expectedLength) {
|
|
data = pkDecompress(data)
|
|
}
|
|
|
|
return data
|
|
}
|
|
|
|
func decompressMulti(data []byte, expectedLength uint32) []byte {
|
|
copmressionType := data[0]
|
|
switch copmressionType {
|
|
case 1: // Huffman
|
|
panic("huffman decompression not supported")
|
|
case 2: // ZLib/Deflate
|
|
return deflate(data[1:])
|
|
case 8: // PKLib/Impode
|
|
return pkDecompress(data[1:])
|
|
case 0x10: // BZip2
|
|
panic("bzip2 decompression not supported")
|
|
case 0x80: // IMA ADPCM Stereo
|
|
return Compression.WavDecompress(data[1:], 2)
|
|
//return MpqWavCompression.Decompress(sinput, 2);
|
|
//panic("ima adpcm sterio decompression not supported")
|
|
case 0x40: // IMA ADPCM Mono
|
|
//return MpqWavCompression.Decompress(sinput, 1)
|
|
panic("mpq wav decompression not supported")
|
|
case 0x12:
|
|
panic("lzma decompression not supported")
|
|
// Combos
|
|
case 0x22:
|
|
// TODO: sparse then zlib
|
|
panic("sparse decompression + deflate decompression not supported")
|
|
case 0x30:
|
|
// TODO: sparse then bzip2
|
|
panic("sparse decompression + bzip2 decompression not supported")
|
|
case 0x41:
|
|
sinput := Compression.HuffmanDecompress(data[1:])
|
|
sinput = Compression.WavDecompress(sinput, 1)
|
|
tmp := make([]byte, len(sinput))
|
|
copy(tmp, sinput)
|
|
return tmp
|
|
case 0x48:
|
|
//byte[] result = PKDecompress(sinput, outputLength);
|
|
//return MpqWavCompression.Decompress(new MemoryStream(result), 1);
|
|
panic("pk + mpqwav decompression not supported")
|
|
case 0x81:
|
|
sinput := Compression.HuffmanDecompress(data[1:])
|
|
sinput = Compression.WavDecompress(sinput, 2)
|
|
tmp := make([]byte, len(sinput))
|
|
copy(tmp, sinput)
|
|
return tmp
|
|
case 0x88:
|
|
//byte[] result = PKDecompress(sinput, outputLength);
|
|
//return MpqWavCompression.Decompress(new MemoryStream(result), 2);
|
|
panic("pk + wav decompression not supported")
|
|
default:
|
|
panic(fmt.Sprintf("decompression not supported for unknown compression type %X", copmressionType))
|
|
}
|
|
}
|
|
|
|
func deflate(data []byte) []byte {
|
|
b := bytes.NewReader(data)
|
|
r, err := zlib.NewReader(b)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
buffer := new(bytes.Buffer)
|
|
_, err = buffer.ReadFrom(r)
|
|
if err != nil {
|
|
log.Panic(err)
|
|
}
|
|
err = r.Close()
|
|
if err != nil {
|
|
log.Panic(err)
|
|
}
|
|
return buffer.Bytes()
|
|
}
|
|
|
|
func pkDecompress(data []byte) []byte {
|
|
b := bytes.NewReader(data)
|
|
r, err := blast.NewReader(b)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
buffer := new(bytes.Buffer)
|
|
_, err = buffer.ReadFrom(r)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
err = r.Close()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return buffer.Bytes()
|
|
}
|