Implement content encryption

This additionally allows a way to convert a WAD (struct) to a genuine WAD file.

Author: spotlightishere

file.go

@@ -62,6 +62,16 @@ func (w *WAD) LoadData(data []byte) error {
 	return nil
 }
 
+func (w *WAD) GetData() []byte {
+	var data []byte
+	for _, content := range w.Data {
+		// Data internally is aligned by 64 bytes.
+		data = append(data, pad(content.RawData)...)
+	}
+
+	return data
+}
+
 func (d *WADFile) DecryptData(titleKey [16]byte) error {
 	content := d.ContentRecord
 
@@ -101,3 +111,49 @@ func (d *WADFile) DecryptData(titleKey [16]byte) error {
 	d.RawData = decryptedData
 	return nil
 }
+
+func (d *WADFile) EncryptData(titleKey [16]byte) error {
+	content := d.ContentRecord
+
+	// The title's decrypted key will be what we'll encrypt with.
+	block, err := aes.NewCipher(titleKey[:])
+	if err != nil {
+		return err
+	}
+
+	// The IV we'll use will be the two bytes sourced from the content's index,
+	// padded with 14 null bytes.
+	var indexBytes [2]byte
+	binary.BigEndian.PutUint16(indexBytes[:], content.Index)
+
+	iv := make([]byte, 16)
+	iv[0] = indexBytes[0]
+	iv[1] = indexBytes[1]
+
+	blockMode := cipher.NewCBCEncrypter(block, iv)
+
+	// One must encrypt content to 16 bytes.
+	// We pad with null bytes.
+	paddedSize := uint32(content.Size)
+	leftover := paddedSize % 16
+	if leftover != 0 {
+		paddedSize += 16 - leftover
+	}
+
+	decryptedData := make([]byte, paddedSize)
+	copy(decryptedData, d.RawData)
+
+	// The resulting encrypted contents is the same size as our adjusted input, including padding.
+	encryptedData := make([]byte, len(decryptedData))
+
+	// ...and we're off!
+	blockMode.CryptBlocks(encryptedData, decryptedData)
+
+	// Update the content record to reflect the hash of our origin content.
+	sha := sha1.Sum(d.RawData)
+	d.Hash = sha
+
+	// We're all set!
+	d.RawData = encryptedData
+	return nil
+}

header.go

@@ -31,7 +31,7 @@ func (w *WAD) LoadHeader(source []byte) error {
 
 func (w *WAD) GetHeader() ([]byte, error) {
 	var tmp bytes.Buffer
-	err := binary.Write(&tmp, binary.BigEndian, w)
+	err := binary.Write(&tmp, binary.BigEndian, w.Header)
 	if err != nil {
 		panic(err)
 	}

wad.go

@@ -34,6 +34,14 @@ func getPadding(size uint32) uint32 {
 	}
 }
 
+func pad(src []byte) []byte {
+	length := (uint32)(len(src))
+	paddedSize := getPadding(length)
+	resized := make([]byte, paddedSize+length)
+	copy(resized, src)
+	return resized
+}
+
 // You use readable when you want to stagger contents and not use scary splice related methods.
 type readable struct {
 	data       []byte
@@ -118,3 +126,46 @@ func LoadWAD(contents []byte) (*WAD, error) {
 
 	return &wad, nil
 }
+
+func (w *WAD) GetWAD(wadType WADType) ([]byte, error) {
+	tmd, err  := w.GetTMD()
+	if err != nil {
+		return nil, err
+	}
+
+	ticket, err  := w.GetTicket()
+	if err != nil {
+		return nil, err
+	}
+
+	data  := w.GetData()
+
+	// Create a header with our sourced content.
+	header := WADHeader{
+		// This size is fixed, and matches Nintendo's size.
+		HeaderSize:      0x20,
+		WADType:         wadType,
+		CertificateSize: (uint32)(len(w.CertificateChain)),
+		CRLSize:         (uint32)(len(w.CertificateRevocationList)),
+		TicketSize:      (uint32)(len(ticket)),
+		TMDSize:         (uint32)(len(tmd)),
+		DataSize:        (uint32)(len(data)),
+		MetaSize:        (uint32)(len(w.Meta)),
+	}
+	w.Header = header
+
+	headerContents, err := w.GetHeader()
+	if err != nil {
+		return nil, err
+	}
+
+	var final []byte
+	final = append(final, pad(headerContents)...)
+	final = append(final, pad(w.CertificateChain)...)
+	final = append(final, pad(w.CertificateRevocationList)...)
+	final = append(final, pad(ticket)...)
+	final = append(final, pad(tmd)...)
+	final = append(final, pad(data)...)
+	final = append(final, pad(w.Meta)...)
+	return final, nil
+}