From 7f32319e2e335353e57d7a3de4efb3ccbcffe5e7 Mon Sep 17 00:00:00 2001
From: nilaoda <nilaoda@live.com>
Date: Tue, 26 Jul 2022 00:49:54 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=8D=95=E6=96=87=E4=BB=B6?=
 =?UTF-8?q?=E6=97=A0=E6=B3=95=E6=AD=A3=E7=A1=AE=E8=A7=A3=E5=AF=86=E7=9A=84?=
 =?UTF-8?q?=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../DownloadManager/SimpleDownloadManager.cs  | 33 +++++++++++--------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/src/N_m3u8DL-RE/DownloadManager/SimpleDownloadManager.cs b/src/N_m3u8DL-RE/DownloadManager/SimpleDownloadManager.cs
index 061939b..f4f7381 100644
--- a/src/N_m3u8DL-RE/DownloadManager/SimpleDownloadManager.cs
+++ b/src/N_m3u8DL-RE/DownloadManager/SimpleDownloadManager.cs
@@ -2,20 +2,13 @@
 using N_m3u8DL_RE.Common.Entity;
 using N_m3u8DL_RE.Common.Log;
 using N_m3u8DL_RE.Common.Resource;
-using N_m3u8DL_RE.Common.Util;
 using N_m3u8DL_RE.Config;
 using N_m3u8DL_RE.Downloader;
 using N_m3u8DL_RE.Entity;
 using N_m3u8DL_RE.Util;
 using Spectre.Console;
-using System;
 using System.Collections.Concurrent;
-using System.Collections.Generic;
-using System.Diagnostics;
-using System.IO;
-using System.Linq;
 using System.Text;
-using System.Threading.Tasks;
 
 namespace N_m3u8DL_RE.DownloadManager
 {
@@ -34,6 +27,15 @@ namespace N_m3u8DL_RE.DownloadManager
 
         private async Task<bool> DownloadStreamAsync(StreamSpec streamSpec, ProgressTask task)
         {
+            string? ReadInit(byte[] data)
+            {
+                var info = MP4InitUtil.ReadInit(data);
+                if (info.Scheme != null) Logger.WarnMarkUp($"[grey]Type: {info.Scheme}[/]");
+                if (info.PSSH != null) Logger.WarnMarkUp($"[grey]PSSH(WV): {info.PSSH}[/]");
+                if (info.KID != null) Logger.WarnMarkUp($"[grey]KID: {info.KID}[/]");
+                return info.KID;
+            }
+
             ConcurrentDictionary<MediaSegment, DownloadResult?> FileDic = new();
 
             var segments = streamSpec.Playlist?.MediaParts.SelectMany(m => m.MediaSegments);
@@ -98,11 +100,7 @@ namespace N_m3u8DL_RE.DownloadManager
                 if (result != null && result.Success) 
                 {
                     var data = File.ReadAllBytes(result.ActualFilePath);
-                    var info = MP4InitUtil.ReadInit(data);
-                    if (info.Scheme != null) Logger.WarnMarkUp($"[grey]Type: {info.Scheme}[/]");
-                    if (info.PSSH != null) Logger.WarnMarkUp($"[grey]PSSH(WV): {info.PSSH}[/]");
-                    if (info.KID != null) Logger.WarnMarkUp($"[grey]KID: {info.KID}[/]");
-                    currentKID = info.KID;
+                    currentKID = ReadInit(data);
                     //实时解密
                     if (DownloaderConfig.MP4RealTimeDecryption && streamSpec.Playlist.MediaInit.EncryptInfo.Method != Common.Enum.EncryptMethod.NONE)
                     {
@@ -363,8 +361,17 @@ namespace N_m3u8DL_RE.DownloadManager
             //调用mp4decrypt解密
             if (File.Exists(output) && !DownloaderConfig.MP4RealTimeDecryption && DownloaderConfig.Keys != null && DownloaderConfig.Keys.Length > 0) 
             {
-                if (totalCount > 1 && streamSpec.Playlist!.MediaParts.First().MediaSegments.First().EncryptInfo.Method != Common.Enum.EncryptMethod.NONE) 
+                if (totalCount >= 1 && streamSpec.Playlist!.MediaParts.First().MediaSegments.First().EncryptInfo.Method != Common.Enum.EncryptMethod.NONE) 
                 {
+                    if (string.IsNullOrEmpty(currentKID))
+                    {
+                        using (var fs = File.OpenRead(output))
+                        {
+                            var header = new byte[4096]; //4KB
+                            fs.Read(header);
+                            currentKID = ReadInit(header);
+                        }
+                    }
                     var enc = output;
                     var dec = Path.Combine(Path.GetDirectoryName(enc)!, Path.GetFileNameWithoutExtension(enc) + "_dec" + Path.GetExtension(enc));
                     Logger.InfoMarkUp($"[grey]Decrypting...[/]");