From e0f184bc05a62a34e73a235375ba893db2108c92 Mon Sep 17 00:00:00 2001
From: Emmm Monster <58943012+emmmx@users.noreply.github.com>
Date: Sat, 26 Dec 2020 04:14:41 +0800
Subject: [PATCH] Add KWM Decoder Improve QMC Decoder

---
 algo/kwm/kwm.go         | 125 ++++++++++++++++++++++++++++++++++++++++
 algo/qmc/mask_key256.go |   2 +-
 2 files changed, 126 insertions(+), 1 deletion(-)
 create mode 100644 algo/kwm/kwm.go

diff --git a/algo/kwm/kwm.go b/algo/kwm/kwm.go
new file mode 100644
index 0000000..b057cf8
--- /dev/null
+++ b/algo/kwm/kwm.go
@@ -0,0 +1,125 @@
+package kwm
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"github.com/davecgh/go-spew/spew"
+	"github.com/umlock-music/cli/algo/common"
+	"strconv"
+	"strings"
+	"unicode"
+)
+
+var (
+	magicHeader = []byte{
+		0x79, 0x65, 0x65, 0x6C, 0x69, 0x6F, 0x6E, 0x2D,
+		0x6B, 0x75, 0x77, 0x6F, 0x2D, 0x74, 0x6D, 0x65}
+	ErrKwFileSize    = errors.New("kwm invalid file size")
+	ErrKwMagicHeader = errors.New("kwm magic header not matched")
+)
+
+const keyPreDefined = "MoOtOiTvINGwd2E6n0E1i7L5t2IoOoNk"
+
+type Decoder struct {
+	file []byte
+
+	key       []byte
+	outputExt string
+	bitrate   int
+	mask      []byte
+
+	audio []byte
+}
+
+func (d *Decoder) GetCoverImage() []byte {
+	return nil
+}
+
+func (d *Decoder) GetAudioData() []byte {
+	return d.audio
+}
+
+func (d *Decoder) GetAudioExt() string {
+	return d.outputExt
+}
+
+func (d *Decoder) GetMeta() common.Meta {
+	return nil
+}
+
+func NewDecoder(data []byte) *Decoder {
+	//todo: Notice the input data will be changed for now
+	return &Decoder{file: data}
+}
+
+func (d *Decoder) Validate() error {
+	lenData := len(d.file)
+	if lenData < 1024 {
+		return ErrKwFileSize
+	}
+	if !bytes.Equal(magicHeader, d.file[:16]) {
+		return ErrKwMagicHeader
+	}
+
+	return nil
+}
+
+func generateMask(key []byte) []byte {
+	keyInt := binary.LittleEndian.Uint64(key)
+	keyStr := strconv.FormatUint(keyInt, 10)
+	keyStrTrim := padOrTruncate(keyStr, 32)
+	mask := make([]byte, 32)
+	for i := 0; i < 32; i++ {
+		mask[i] = keyPreDefined[i] ^ keyStrTrim[i]
+	}
+	return mask
+}
+
+func (d *Decoder) parseBitrateAndType() {
+	bitType := string(bytes.TrimRight(d.file[0x30:0x38], string(byte(0))))
+	charPos := 0
+	for charPos = range bitType {
+		if !unicode.IsNumber(rune(bitType[charPos])) {
+			break
+		}
+	}
+	var err error
+	d.bitrate, err = strconv.Atoi(bitType[:charPos])
+	if err != nil {
+		d.bitrate = 0
+	}
+	d.outputExt = strings.ToLower(bitType[charPos:])
+
+}
+
+func (d *Decoder) Decode() error {
+	d.parseBitrateAndType()
+
+	d.mask = generateMask(d.file[0x18:0x20])
+
+	d.audio = d.file[1024:]
+	dataLen := len(d.audio)
+	spew.Dump(d.audio[:1024])
+	for i := 0; i < dataLen; i++ {
+		d.audio[i] ^= d.mask[i&0x1F] //equals: [i % 32]
+	}
+	return nil
+}
+
+func padOrTruncate(raw string, length int) string {
+	lenRaw := len(raw)
+	out := raw
+	if lenRaw == 0 {
+		out = string(make([]byte, length))
+	} else if lenRaw > length {
+		out = raw[:length]
+	} else if lenRaw < length {
+		_tmp := make([]byte, 32)
+		for i := 0; i < 32; i++ {
+			_tmp[i] = raw[i%lenRaw]
+		}
+		out = string(_tmp)
+	}
+	return out
+}
diff --git a/algo/qmc/mask_key256.go b/algo/qmc/mask_key256.go
index bb50393..11f91cb 100644
--- a/algo/qmc/mask_key256.go
+++ b/algo/qmc/mask_key256.go
@@ -143,7 +143,7 @@ func detectMgg256Mask(input []byte) (*Key256Mask, error) {
 		if confidence > 0 {
 			mask := input[idx128] ^ spHeader[idx128]
 
-			idx44 := key256Mapping128to44[idx128%128]
+			idx44 := key256Mapping128to44[idx128&0x7f] // equals: [idx128 % 128]
 			if _, ok2 := matrixConf[idx44][mask]; ok2 {
 				matrixConf[idx44][mask] += confidence
 			} else {