From e753b0da9b513aeb60c30ca5fe8b56238754d309 Mon Sep 17 00:00:00 2001
From: CastagnaIT <gottardo.stefano.83@gmail.com>
Date: Tue, 27 Aug 2024 18:19:10 +0200
Subject: [PATCH] test

---
 src/Session.cpp            | 18 +++++++++++++++++
 src/decrypters/Helpers.cpp | 40 ++++++++++++++++++++++++++++++++++++++
 src/decrypters/Helpers.h   |  1 +
 src/parser/DASHTree.cpp    |  2 +-
 4 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/src/Session.cpp b/src/Session.cpp
index 8012a4358..989aa18c0 100644
--- a/src/Session.cpp
+++ b/src/Session.cpp
@@ -17,6 +17,7 @@
 #include "common/Chooser.h"
 #include "decrypters/DrmFactory.h"
 #include "decrypters/Helpers.h"
+#include "parser/PRProtectionParser.h"
 #include "utils/Base64Utils.h"
 #include "utils/CurlUtils.h"
 #include "utils/StringUtils.h"
@@ -451,6 +452,23 @@ bool CSession::InitializeDRM(bool addDefaultKID /* = false */)
         initData = BASE64::Decode(licenseDataStr);
       }
 
+      if (!initData.empty() && defaultKidStr.empty())
+      {
+        CPsshParser parser;
+        if (parser.Parse(initData))
+        {
+          if (!parser.GetData().empty())
+          {
+            auto data = DRM::GetKIDWidevinePsshData(parser.GetData());
+            if (!data.empty())
+            {
+              LOG::LogF(LOGDEBUG, "KID PARSED FROM PSSH");
+              defaultKidStr = STRING::ToHexadecimal(data);
+            }
+          }
+        }
+      }
+
       //! @todo: as is implemented InitializeDRM will initialize all PSSHSet's also when are not used,
       //!   therefore ExtractStreamProtectionData can perform many (not needed) downloads of mp4 init files
       if ((initData.empty() && licenseType != DRM::KS_CLEARKEY) || defaultKidStr.empty())
diff --git a/src/decrypters/Helpers.cpp b/src/decrypters/Helpers.cpp
index 02d226362..a089e7f28 100644
--- a/src/decrypters/Helpers.cpp
+++ b/src/decrypters/Helpers.cpp
@@ -61,6 +61,21 @@ void WriteProtobufVarint(std::vector<uint8_t>& data, int size)
   } while (size > 0);
 }
 
+int ReadProtobufVarint(const std::vector<uint8_t>& data, size_t& offset)
+{
+  int value = 0;
+  int shift = 0;
+  while (true)
+  {
+    uint8_t byte = data[offset++];
+    value |= (byte & 0x7F) << shift;
+    if (!(byte & 0x80))
+      break;
+    shift += 7;
+  }
+  return value;
+}
+
 /*!
  * \brief Replace in a vector, a sequence of vector data with another one.
  * \param data The data to be modified
@@ -229,6 +244,31 @@ bool DRM::IsValidPsshHeader(const std::vector<uint8_t>& pssh)
   return pssh.size() >= 8 && std::equal(pssh.begin() + 4, pssh.begin() + 8, PSSHBOX_HEADER_PSSH);
 }
 
+std::vector<uint8_t> DRM::GetKIDWidevinePsshData(const std::vector<uint8_t>& wvPsshData)
+{
+  size_t offset = 0;
+  while (offset < wvPsshData.size())
+  {
+    uint8_t tag = wvPsshData[offset++];
+    int fieldNumber = tag >> 3;
+    int wireType = tag & 0x07;
+
+    if (fieldNumber == 2 && wireType == 2) // "key_id" field, id: 2
+    {
+      int length = ReadProtobufVarint(wvPsshData, offset);
+      std::vector<uint8_t> kid(wvPsshData.begin() + offset, wvPsshData.begin() + offset + length);
+      return kid;
+    }
+    else // Skip field
+    {
+      int length = ReadProtobufVarint(wvPsshData, offset);
+      if (wireType != 0)
+        offset += length;
+    }
+  }
+  return {}; // Not found
+}
+
 bool DRM::MakeWidevinePsshData(const std::vector<uint8_t>& kid,
                                std::vector<uint8_t> contentIdData,
                                std::vector<uint8_t>& wvPsshData)
diff --git a/src/decrypters/Helpers.h b/src/decrypters/Helpers.h
index e6080e2ec..3d2d91c72 100644
--- a/src/decrypters/Helpers.h
+++ b/src/decrypters/Helpers.h
@@ -81,6 +81,7 @@ std::vector<uint8_t> ConvertKidToUUIDVec(const std::vector<uint8_t>& kid);
  * \param kid The PlayReady KeyId
  * \return The Widevine KeyId, otherwise empty if fails.
  */
+std::vector<uint8_t> GetKIDWidevinePsshData(const std::vector<uint8_t>& wvPsshData);
 std::vector<uint8_t> ConvertPrKidtoWvKid(std::vector<uint8_t> kid);
 
 bool IsValidPsshHeader(const std::vector<uint8_t>& pssh);
diff --git a/src/parser/DASHTree.cpp b/src/parser/DASHTree.cpp
index a0cd6b022..ab21762ae 100644
--- a/src/parser/DASHTree.cpp
+++ b/src/parser/DASHTree.cpp
@@ -1365,7 +1365,7 @@ bool adaptive::CDashTree::GetProtectionData(
 
   // There are no constraints on the Kid format, it is recommended to be as UUID but not mandatory
   STRING::ReplaceAll(selectedKid, "-", "");
-  kid = selectedKid;
+  //kid = selectedKid;
 
   return isEncrypted;
 }