gabubellon · gabubellon · Oct 12, 2021 · Oct 12, 2021
diff --git a/streamyard_down/__main__.py b/streamyard_down/__main__.py
@@ -6,6 +6,6 @@ def main():
     streamyardown = StreamYardDownload(**args.get_args_dict())
     streamyardown.start_download()
 
+
 if __name__ == "__main__":
     main()
-
diff --git a/streamyard_down/args.py b/streamyard_down/args.py
@@ -2,25 +2,30 @@
 import calendar
 from datetime import date, datetime
 
-parser = argparse.ArgumentParser(description="Download finished streamns from StreamYard")
+parser = argparse.ArgumentParser(
+    description="Download finished streamns from StreamYard"
+)
+
+streamyard_group = parser.add_argument_group("streamyard")
+s3_group = parser.add_argument_group("s3")
 
-parser.add_argument(
+streamyard_group.add_argument(
     "--email",
     type=str,
     help="registred email on streamyard ",
     dest="email",
     required=True,
 )
 
-parser.add_argument(
+streamyard_group.add_argument(
     "-lc",
     "--list_choise",
     action="store_true",
     help="Show a list of streamns to dowload",
     dest="list_choise",
 )
 
-parser.add_argument(
+streamyard_group.add_argument(
     "-p",
     "--path",
     type=str,
@@ -29,7 +34,7 @@
     dest="path",
 )
 
-parser.add_argument(
+streamyard_group.add_argument(
     "-s",
     "--start_date",
     type=date,
@@ -38,7 +43,7 @@
     dest="start_date",
 )
 
-parser.add_argument(
+streamyard_group.add_argument(
     "-e",
     "--end_date",
     type=date,
@@ -54,15 +59,15 @@
     dest="end_date",
 )
 
-parser.add_argument(
+streamyard_group.add_argument(
     "-nl",
     "--new_login",
     action="store_true",
     help="Force a new login and ignore cookie cache (default: False)",
     dest="new_login",
 )
 
-parser.add_argument(
+streamyard_group.add_argument(
     "-c",
     "--chuck_size",
     type=int,
@@ -71,7 +76,7 @@
     dest="chuck_size",
 )
 
-parser.add_argument(
+streamyard_group.add_argument(
     "-t",
     "--threads",
     type=int,
@@ -80,16 +85,47 @@
     dest="threads",
 )
 
-parser.add_argument(
+s3_group.add_argument(
     "-s3",
     "--upload_s3",
     action="store_true",
     help="automatic upload files to a s3 bucket (defaul: False)",
     dest="upload",
 )
 
+s3_group.add_argument(
+    "--bucket",
+    type=str,
+    help="s3 bucket to save files on formart s3://bucket-name",
+    dest="bucket",
+)
+
+s3_group.add_argument(
+    "--prefix",
+    type=str,
+    help="s3 prefix to save file on bucket",
+    dest="prefix",
+)
+
 args = parser.parse_args()
 
 
 def get_args_dict():
-    return vars(args)
+
+    streamyard_keys = [
+        "email",
+        "list_choise",
+        "path",
+        "start_date",
+        "end_date",
+        "new_login",
+        "chuck_size",
+        "threads",
+        "upload",
+    ]
+
+    streamyard_args = {}
+    for arg in streamyard_keys:
+        streamyard_args[arg] = vars(args)[arg]
+
+    return streamyard_args
diff --git a/streamyard_down/s3.py b/streamyard_down/s3.py
@@ -0,0 +1,19 @@
+import os
+
+import boto3
+from loguru import logger
+
+# TODO Transform in a class
+
+s3_prefix = ""
+s3_bucket = ""
+
+
+def send_to_s3(sent_file):
+    key = "{}/{}".format(s3_prefix, os.path.basename(sent_file))
+    logger.info(f"Sending {sent_file} to bucket {s3_bucket}/{key}")
+
+    s3_client = boto3.client("s3")
+    s3_client.upload_file(sent_file, s3_bucket, key)
+
+    os.remove(sent_file)
diff --git a/streamyard_down/streamyard.py b/streamyard_down/streamyard.py
@@ -6,7 +6,6 @@
 from concurrent import futures
 from concurrent.futures import ProcessPoolExecutor
 
-import boto3
 import pandas as pd
 import requests
 from loguru import logger
@@ -37,6 +36,7 @@ def __init__(
         self.end_date = end_date
         self.request_session = self.create_session()
         self.list_choise = list_choise
+        self.quit = False
 
     def create_session(self):
         logger.info("Criando Sessão")
@@ -71,8 +71,9 @@ def download_file(self, file_name, request_url):
                         show_log = False
                         previus_done = done
 
+        # TODO implement upload
         if self.upload:
-            self.send_to_s3(os.path.join(self.path, file_name))
+            pass
 
     def get_cookie(self):
         logger.info("Carregando Cookie")
@@ -186,59 +187,61 @@ def dowload(self, broadcast_to_download):
             logger.info(f"Download da stream {file_name} completo ")
 
     def download_broadcast(self, stream_info):
-        stream_id = stream_info.get("stream_id")
-        file_name = stream_info.get("file_name")
-        video_filename = stream_info.get("video_filename")
-        audio_filename = stream_info.get("audio_filename")
-
-        logger.info(f"Download stream id:{stream_id} name:{file_name}")
-
-        # ESSA CHAMADA NÃO FUNCIONA PELA API, APENAS ENTRANDO NO SITE E CLICANDO NO BOTÃO
-        get_url = self.request_session.post(
-            cfg.CREATE_DOWNLOADS_URL.format(stream_id=stream_id),
-            data=dict(csrfToken=self.LOGGED_TOKEN),
-            headers=dict(Referer=cfg.BROAD_CAST_URL),
-        )
-
-        logger.info(f"{get_url.text}")
+        while not self.quit:
+            stream_id = stream_info.get("stream_id")
+            file_name = stream_info.get("file_name")
+            video_filename = stream_info.get("video_filename")
+            audio_filename = stream_info.get("audio_filename")
+
+            logger.info(f"Download stream id:{stream_id} name:{file_name}")
+
+            # ESSA CHAMADA NÃO FUNCIONA PELA API, APENAS ENTRANDO NO SITE E CLICANDO NO BOTÃO
+            get_url = self.request_session.post(
+                cfg.CREATE_DOWNLOADS_URL.format(stream_id=stream_id),
+                data=dict(csrfToken=self.LOGGED_TOKEN),
+                headers=dict(Referer=cfg.BROAD_CAST_URL),
+            )
 
-        while True:
-            logger.info(f"Gerando Links de download")
+            logger.info(f"{get_url.text}")
 
-            # COMO A CHAMADA NEM SEMPRE FUNCIONA ESSE REQUEST RETORNA ESSE
-            make_urls = self.request_session.get(
-                cfg.CREATE_DOWNLOADS_URL.format(stream_id=stream_id)
-            )
+            while True:
+                logger.info(f"Gerando Links de download")
 
-            logger.info(f"{make_urls.text}")
-            status = json.loads(make_urls.text).get("status")
+                # COMO A CHAMADA NEM SEMPRE FUNCIONA ESSE REQUEST RETORNA ESSE
+                make_urls = self.request_session.get(
+                    cfg.CREATE_DOWNLOADS_URL.format(stream_id=stream_id)
+                )
 
-            logger.info(f"Status: {status}")
+                logger.info(f"{make_urls.text}")
+                status = json.loads(make_urls.text).get("status")
 
-            if status != "creating":
-                break
+                logger.info(f"Status: {status}")
 
-            time.sleep(10)
+                if status != "creating":
+                    break
 
-        download_url = self.request_session.get(
-            cfg.DOWNLOAD_URL.format(stream_id=stream_id)
-        )
+                time.sleep(10)
 
-        if json.loads(download_url.text).get("audioUrl"):
-            logger.info(f"Download do audio")
-            self.download_file(
-                file_name=audio_filename,
-                request_url=json.loads(download_url.text).get("audioUrl"),
+            download_url = self.request_session.get(
+                cfg.DOWNLOAD_URL.format(stream_id=stream_id)
             )
 
-        if json.loads(download_url.text).get("videoUrl"):
-            logger.info(f"Download do video")
-            self.download_file(
-                file_name=video_filename,
-                request_url=json.loads(download_url.text).get("videoUrl"),
-            )
+            if json.loads(download_url.text).get("audioUrl"):
+                logger.info(f"Download do audio")
+                self.download_file(
+                    file_name=audio_filename,
+                    request_url=json.loads(download_url.text).get("audioUrl"),
+                )
+
+            if json.loads(download_url.text).get("videoUrl"):
+                logger.info(f"Download do video")
+                self.download_file(
+                    file_name=video_filename,
+                    request_url=json.loads(download_url.text).get("videoUrl"),
+                )
 
-        return file_name
+            return file_name
+        return
 
     def start_download(self):
         self.login()
@@ -247,7 +250,7 @@ def start_download(self):
         for index, item in enumerate(download_list):
             options += f"** ID:{index} - STREAM:{item.get('file_name')} - DATE:{item.get('stream_date')}\n"
 
-        to_download=[]
+        to_download = []
         if self.list_choise:
             message = f"""
             ############################
@@ -265,24 +268,18 @@ def start_download(self):
             """
             ids = input(message)
 
-            to_download=[]
-            escolhas=""
-            for id in ids.split(','):
+            to_download = []
+            escolhas = ""
+            for id in ids.split(","):
 
                 data = download_list[int(id)]
                 escolhas += f"** ID:{id} - STREAM:{data.get('file_name')} - DATE:{data.get('stream_date')}\n"
                 to_download.append(data)
             print(f"OS SEGUINTES IDs serão baixados {ids}\n{escolhas}")
         else:
             to_download = download_list
-
-        self.dowload(to_download)
-
-    def send_to_s3(self, sent_file):
-        key = "{}/{}".format(cfg.S3_PREFIX, os.path.basename(sent_file))
-        logger.info(f"Sending {sent_file} to bucket s3://{cfg.S3_BUCKET}/{key}")
-
-        s3_client = boto3.client("s3")
-        s3_client.upload_file(sent_file, cfg.S3_BUCKET, key)
 
-        os.remove(sent_file)
+        try:
+            self.dowload(to_download)
+        except KeyboardInterrupt:
+            self.quit = True