From 5e1a07e793f12045a5e2d300ac789cab2a5e7610 Mon Sep 17 00:00:00 2001
From: samedamci <samedamci@disroot.org>
Date: Sat, 11 Sep 2021 21:06:48 +0200
Subject: [PATCH] Add new bool option to get_stream method

---
 mattebox/app.py   | 8 ++++++--
 mattebox/types.py | 6 ++----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/mattebox/app.py b/mattebox/app.py
index c6afdbc..603e9d8 100644
--- a/mattebox/app.py
+++ b/mattebox/app.py
@@ -55,8 +55,12 @@ class MatteBOX:
         recordings = [Program.from_recording(r) for r in res["entities"]]
         return recordings
 
-    def get_stream(self, program: Program) -> str:
+    def get_stream(self, program: Program, end: bool = True) -> str:
         service_type = "TIMESHIFT_TV" if program.type == "program" else "NPVR"
+        if end:
+            ts_end = program.ts_stop
+        else:
+            ts_end = now_timestamp() - 60000 * 3  # 3 minutes ago
 
         params = {
             "channelKey": program.channel,
@@ -65,7 +69,7 @@ class MatteBOX:
             "subscriptionCode": self.subscription_code,
             "deviceType": "STB",
             "fromTimestamp": program.ts_start,
-            "toTimestamp": program.ts_stop,
+            "toTimestamp": ts_end,
         }
 
         res = self.__get("/sws/server/streaming/uris.json", params=params)
diff --git a/mattebox/types.py b/mattebox/types.py
index 5db40ee..a5adcd9 100644
--- a/mattebox/types.py
+++ b/mattebox/types.py
@@ -1,5 +1,4 @@
 from __future__ import annotations
-from .helpers import now_timestamp
 
 
 class Program:
@@ -8,7 +7,7 @@ class Program:
     content_id: str
     description: str
     ts_start: int
-    ts_stop: int
+    ts_stop: int = None
     type: str
 
     @classmethod
@@ -20,7 +19,7 @@ class Program:
             "content_id": data["epgId"],
             "description": data["shortDescription"],
             "ts_start": data["startTimestamp"],
-            "ts_stop": now_timestamp() - 60000 * 3,  # 3 minutes ago
+            "ts_stop": data["endTimestamp"],
             "type": "program",
         }
         return obj
@@ -34,7 +33,6 @@ class Program:
             "content_id": data["pvrProgramId"],
             "description": data["longDescription"],
             "ts_start": data["startTime"],
-            "ts_stop": now_timestamp(),
             "type": "recording",
         }
         return obj