DiscIO: Implement re-encryption of Wii partition data

Source/Core/DiscIO/DirectoryBlob.h

@@ -4,6 +4,7 @@
 
 #pragma once
 
+#include <array>
 #include <cstddef>
 #include <map>
 #include <memory>
@@ -16,6 +17,7 @@
 #include "Common/CommonTypes.h"
 #include "Common/FileUtil.h"
 #include "DiscIO/Blob.h"
+#include "DiscIO/WiiEncryptionCache.h"
 
 namespace File
 {
@@ -27,16 +29,23 @@ namespace DiscIO
 {
 enum class PartitionType : u32;
 
+class DirectoryBlobReader;
+
 // Returns true if the path is inside a DirectoryBlob and doesn't represent the DirectoryBlob itself
 bool ShouldHideFromGameList(const std::string& volume_path);
 
 class DiscContent
 {
 public:
-  using ContentSource = std::variant<std::string, const u8*>;
+  using ContentSource =
+      std::variant<std::string,          // File
+                   const u8*,            // Memory
+                   DirectoryBlobReader*  // Partition (which one it is is determined by m_offset)
+                   >;
 
   DiscContent(u64 offset, u64 size, const std::string& path);
   DiscContent(u64 offset, u64 size, const u8* data);
+  DiscContent(u64 offset, u64 size, DirectoryBlobReader* blob);
 
   // Provided because it's convenient when searching for DiscContent in an std::set
   explicit DiscContent(u64 offset);
@@ -69,6 +78,7 @@ public:
   }
   void Add(u64 offset, u64 size, const std::string& path);
   void Add(u64 offset, u64 size, const u8* data);
+  void Add(u64 offset, u64 size, DirectoryBlobReader* blob);
   u64 CheckSizeAndAdd(u64 offset, const std::string& path);
   u64 CheckSizeAndAdd(u64 offset, u64 max_size, const std::string& path);
 
@@ -96,6 +106,9 @@ public:
   const std::vector<u8>& GetHeader() const { return m_disc_header; }
   const DiscContentContainer& GetContents() const { return m_contents; }
 
+  const std::array<u8, VolumeWii::AES_KEY_SIZE>& GetKey() const { return m_key; }
+  void SetKey(std::array<u8, VolumeWii::AES_KEY_SIZE> key) { m_key = key; }
+
 private:
   void SetDiscHeaderAndDiscType(std::optional<bool> is_wii);
   void SetBI2();
@@ -120,6 +133,8 @@ private:
   std::vector<u8> m_apploader;
   std::vector<u8> m_fst_data;
 
+  std::array<u8, VolumeWii::AES_KEY_SIZE> m_key;
+
   std::string m_root_directory;
   bool m_is_wii = false;
   // GameCube has no shift, Wii has 2 bit shift
@@ -130,6 +145,8 @@ private:
 
 class DirectoryBlobReader : public BlobReader
 {
+  friend DiscContent;
+
 public:
   static std::unique_ptr<DirectoryBlobReader> Create(const std::string& dol_path);
 
@@ -163,11 +180,14 @@ private:
   explicit DirectoryBlobReader(const std::string& game_partition_root,
                                const std::string& true_root);
 
+  bool EncryptPartitionData(u64 offset, u64 size, u8* buffer, u64 partition_data_offset,
+                            u64 partition_data_decrypted_size);
+
   void SetNonpartitionDiscHeader(const std::vector<u8>& partition_header,
                                  const std::string& game_partition_root);
   void SetWiiRegionData(const std::string& game_partition_root);
   void SetPartitions(std::vector<PartitionWithType>&& partitions);
-  void SetPartitionHeader(const DirectoryBlobPartition& partition, u64 partition_address);
+  void SetPartitionHeader(DirectoryBlobPartition* partition, u64 partition_address);
 
   // For GameCube:
   DirectoryBlobPartition m_gamecube_pseudopartition;
@@ -175,8 +195,10 @@ private:
   // For Wii:
   DiscContentContainer m_nonpartition_contents;
   std::map<u64, DirectoryBlobPartition> m_partitions;
+  WiiEncryptionCache m_encryption_cache;
 
   bool m_is_wii;
+  bool m_encrypted;
 
   std::vector<u8> m_disc_header_nonpartition;
   std::vector<u8> m_partition_table;