From c18fe5706b37c385ed4f3eb999ed3fa280dfd1d9 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Sat, 10 Feb 2024 17:48:26 +0000
Subject: [PATCH] Pass DRM and Segment objects to Track OnDecrypted event

---
 devine/commands/dl.py         | 2 +-
 devine/core/manifests/dash.py | 2 +-
 devine/core/manifests/hls.py  | 2 +-
 devine/core/tracks/track.py   | 5 +++--
 4 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/devine/commands/dl.py b/devine/commands/dl.py
index d585c76..8d33417 100644
--- a/devine/commands/dl.py
+++ b/devine/commands/dl.py
@@ -899,7 +899,7 @@ class dl:
                             drm.decrypt(save_path)
                             track.drm = None
                             if callable(track.OnDecrypted):
-                                track.OnDecrypted()
+                                track.OnDecrypted(drm)
                             progress(downloaded="Decrypted", completed=100)
 
                         if isinstance(track, Subtitle):
diff --git a/devine/core/manifests/dash.py b/devine/core/manifests/dash.py
index 129cd88..9d0946a 100644
--- a/devine/core/manifests/dash.py
+++ b/devine/core/manifests/dash.py
@@ -507,7 +507,7 @@ class DASH:
             drm.decrypt(save_path)
             track.drm = None
             if callable(track.OnDecrypted):
-                track.OnDecrypted()
+                track.OnDecrypted(drm)
             progress(downloaded="Decrypted", completed=100)
 
         track.path = save_path
diff --git a/devine/core/manifests/hls.py b/devine/core/manifests/hls.py
index 1639eb6..3a35e31 100644
--- a/devine/core/manifests/hls.py
+++ b/devine/core/manifests/hls.py
@@ -511,7 +511,7 @@ class HLS:
             newest_segment_key[0].decrypt(out_path)
             track.drm = None
             if callable(track.OnDecrypted):
-                track.OnDecrypted()
+                track.OnDecrypted(newest_segment_key[0], segment)
 
         return download_size
 
diff --git a/devine/core/tracks/track.py b/devine/core/tracks/track.py
index e4018fa..36d3f6f 100644
--- a/devine/core/tracks/track.py
+++ b/devine/core/tracks/track.py
@@ -7,6 +7,7 @@ from pathlib import Path
 from typing import Any, Callable, Iterable, Optional, Union
 from uuid import UUID
 
+import m3u8
 import requests
 from langcodes import Language
 
@@ -57,8 +58,8 @@ class Track:
         # TODO: This should realistically be before decryption
         # Called after the Track has been fully downloaded and decrypted
         self.OnDownloaded: Optional[Callable] = None
-        # Called after the Track or a Segment has been decrypted
-        self.OnDecrypted: Optional[Callable] = None
+        # Called after the Track or one of its segments have been decrypted
+        self.OnDecrypted: Optional[Callable[[DRM_T, Optional[m3u8.Segment]], None]] = None
         # Called after the Track has been repackaged
         self.OnRepacked: Optional[Callable] = None
         # Called before the Track is multiplexed