From 502e05dc2bf222a9f0dfe33b8506b5040e4486ff Mon Sep 17 00:00:00 2001
From: Pengfei <zhupf321@gmail.com>
Date: Thu, 29 Jul 2021 16:38:41 +0800
Subject: [PATCH] Add CIABuildType to CIABuilder; Add legit ticket handling

---
 src/core/importer.cpp         |  2 +-
 src/core/ncch/cia_builder.cpp | 38 ++++++++++++++++++++++++++++++-----
 src/core/ncch/cia_builder.h   | 13 ++++++++++--
 3 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/src/core/importer.cpp b/src/core/importer.cpp
index 09e2878..5446851 100644
--- a/src/core/importer.cpp
+++ b/src/core/importer.cpp
@@ -693,7 +693,7 @@ bool SDMCImporter::BuildCIA(const ContentSpecifier& specifier, std::string desti
         }
     }
 
-    const bool ret = cia_builder->Init(destination, tmd, config,
+    const bool ret = cia_builder->Init(CIABuildType::Standard, destination, tmd, config,
                                        FileUtil::GetDirectoryTreeSize(physical_path), callback);
     if (!ret) {
         FileUtil::Delete(destination);
diff --git a/src/core/ncch/cia_builder.cpp b/src/core/ncch/cia_builder.cpp
index 5b45190..1500d22 100644
--- a/src/core/ncch/cia_builder.cpp
+++ b/src/core/ncch/cia_builder.cpp
@@ -48,9 +48,11 @@ private:
 CIABuilder::CIABuilder() = default;
 CIABuilder::~CIABuilder() = default;
 
-bool CIABuilder::Init(const std::string& destination, TitleMetadata tmd_, const Config& config,
-                      std::size_t total_size_, const Common::ProgressCallback& callback_) {
+bool CIABuilder::Init(CIABuildType type_, const std::string& destination, TitleMetadata tmd_,
+                      const Config& config, std::size_t total_size_,
+                      const Common::ProgressCallback& callback_) {
 
+    type = type_;
     header = {};
     meta = {};
 
@@ -119,9 +121,19 @@ bool CIABuilder::WriteCert(const std::string& certs_db_path) {
     return true;
 }
 
-bool CIABuilder::WriteTicket(const std::string& ticket_db_path,
-                             const std::string& enc_title_keys_bin_path) {
-    const auto title_id = tmd.GetTitleID();
+static bool FindLegitTicket(Ticket& out, u64 title_id, const std::string& ticket_db_path) {
+    TicketDB ticket_db(ticket_db_path);
+    if (ticket_db.IsGood() && ticket_db.tickets.count(title_id)) {
+        out = ticket_db.tickets.at(title_id);
+        return true;
+    }
+
+    LOG_ERROR(Core, "Could not find legit ticket for {:016x}", title_id);
+    return false;
+}
+
+static Ticket BuildStandardTicket(u64 title_id, const std::string& ticket_db_path,
+                                  const std::string& enc_title_keys_bin_path) {
     Ticket ticket = BuildFakeTicket(title_id);
 
     // Fill in common_key_index and title_key from either ticket.db (installed tickets)
@@ -143,6 +155,22 @@ bool CIABuilder::WriteTicket(const std::string& ticket_db_path,
     } else {
         LOG_WARNING(Core, "Could not find title key for {:016x}", title_id);
     }
+    return ticket;
+}
+
+bool CIABuilder::WriteTicket(const std::string& ticket_db_path,
+                             const std::string& enc_title_keys_bin_path) {
+
+    const auto title_id = tmd.GetTitleID();
+
+    Ticket ticket;
+    if (type == CIABuildType::Legit) {
+        if (!FindLegitTicket(ticket, title_id, ticket_db_path)) {
+            return false;
+        }
+    } else {
+        ticket = BuildStandardTicket(title_id, ticket_db_path, enc_title_keys_bin_path);
+    }
 
     header.tik_size = ticket.GetSize();
 
diff --git a/src/core/ncch/cia_builder.h b/src/core/ncch/cia_builder.h
index 5b18909..8bc85e3 100644
--- a/src/core/ncch/cia_builder.h
+++ b/src/core/ncch/cia_builder.h
@@ -25,6 +25,12 @@ constexpr std::size_t CIA_METADATA_SIZE = 0x3AC0;
 struct Config;
 class HashedFile;
 
+enum class CIABuildType {
+    Standard,    /// Decrypted CIA with generalized ticket
+    PirateLegit, /// Uses legit TMD and encryption, but with generalized ticket
+    Legit,       /// Fully legit, with personal ticket containing console ID and eshop account
+};
+
 class CIABuilder {
 public:
     explicit CIABuilder();
@@ -34,8 +40,9 @@ public:
      * Initializes the building of the CIA.
      * @return true on success, false otherwise
      */
-    bool Init(const std::string& destination, TitleMetadata tmd, const Config& config,
-              std::size_t total_size, const Common::ProgressCallback& callback);
+    bool Init(CIABuildType type, const std::string& destination, TitleMetadata tmd,
+              const Config& config, std::size_t total_size,
+              const Common::ProgressCallback& callback);
 
     /**
      * Adds an NCCH content to the CIA.
@@ -93,6 +100,8 @@ private:
     bool WriteCert(const std::string& certs_db_path);
     bool WriteTicket(const std::string& ticket_db_path, const std::string& enc_title_keys_bin_path);
 
+    CIABuildType type;
+
     Header header{};
     Metadata meta{};