From 2e7ef6f0083f9b51d2874dcd6a949c43ce006984 Mon Sep 17 00:00:00 2001
From: Stefano Pigozzi <ste.pigozzi@gmail.com>
Date: Sat, 20 Apr 2019 15:26:19 +0200
Subject: [PATCH] Escape filenames to ensure they can be loaded from Discord

---
 royalnet/audio/royalpcmfile.py | 29 +++++++++++++++++++----------
 royalnet/utils/__init__.py     |  3 ++-
 royalnet/utils/safefilename.py |  5 +++++
 3 files changed, 26 insertions(+), 11 deletions(-)
 create mode 100644 royalnet/utils/safefilename.py

diff --git a/royalnet/audio/royalpcmfile.py b/royalnet/audio/royalpcmfile.py
index ef407a45..272b3521 100644
--- a/royalnet/audio/royalpcmfile.py
+++ b/royalnet/audio/royalpcmfile.py
@@ -1,8 +1,13 @@
+import logging
 import ffmpeg
 import os
 import typing
 import time
 from .youtubedl import YtdlFile, YtdlInfo
+from ..utils import safefilename
+
+
+log = logging.getLogger(__name__)
 
 
 class RoyalPCMFile(YtdlFile):
@@ -20,17 +25,21 @@ class RoyalPCMFile(YtdlFile):
             raise FileExistsError("Can't overwrite file")
         # Overwrite the new ytdl_args
         self.ytdl_args = {**self.ytdl_args, **ytdl_args}
-        self.ytdl_args["log"].info(f"Now downloading {info.webpage_url}")
+        log.info(f"Now downloading {info.webpage_url}")
         super().__init__(info, outtmpl=self._ytdl_filename, **self.ytdl_args)
         # Find the audio_filename with a regex (should be video.opus)
-        self.ytdl_args["log"].info(f"Preparing {self.video_filename}...")
+        log.info(f"Preparing {self.video_filename}...")
         # Convert the video to pcm
-        ffmpeg.input(f"./{self.video_filename}") \
-              .output(self.audio_filename, format="s16le", ac=2, ar="48000") \
-              .overwrite_output() \
-              .run(quiet=False)
+        try:
+            ffmpeg.input(f"./{self.video_filename}") \
+                  .output(self.audio_filename, format="s16le", ac=2, ar="48000") \
+                  .overwrite_output() \
+                  .run(quiet=False)
+        except ffmpeg.Error as exc:
+            log.error(f"FFmpeg error: {exc.stderr}")
+            raise
         # Delete the video file
-        self.ytdl_args["log"].info(f"Deleting {self.video_filename}")
+        log.info(f"Deleting {self.video_filename}")
         self.delete_video_file()
 
     def __repr__(self):
@@ -43,12 +52,12 @@ class RoyalPCMFile(YtdlFile):
 
     @property
     def _ytdl_filename(self):
-        return f"./downloads/{self.info.title}-{str(int(self._time))}.ytdl"
+        return f"./downloads/{safefilename(self.info.title)}-{safefilename(str(int(self._time)))}.ytdl"
 
     @property
     def audio_filename(self):
-        return f"./downloads/{self.info.title}-{str(int(self._time))}.pcm"
+        return f"./downloads/{safefilename(self.info.title)}-{safefilename(str(int(self._time)))}.pcm"
 
     def __del__(self):
-        self.ytdl_args["log"].info(f"Deleting {self.audio_filename}")
+        log.info(f"Deleting {self.audio_filename}")
         os.remove(self.audio_filename)
diff --git a/royalnet/utils/__init__.py b/royalnet/utils/__init__.py
index edcb1256..209f6ddb 100644
--- a/royalnet/utils/__init__.py
+++ b/royalnet/utils/__init__.py
@@ -7,6 +7,7 @@ from .classdictjanitor import cdj
 from .sleepuntil import sleep_until
 from .plusformat import plusformat
 from .networkhandler import NetworkHandler
+from .safefilename import safefilename
 
 __all__ = ["asyncify", "Call", "Command", "safeformat", "cdj", "sleep_until", "plusformat", "CommandArgs",
-           "NetworkHandler"]
+           "NetworkHandler", "safefilename"]
diff --git a/royalnet/utils/safefilename.py b/royalnet/utils/safefilename.py
new file mode 100644
index 00000000..3f6d8e2f
--- /dev/null
+++ b/royalnet/utils/safefilename.py
@@ -0,0 +1,5 @@
+import re
+
+
+def safefilename(string: str):
+    return re.sub(r"\W", "_", string)