From 47eebac4d87ae83f22cb025e67faf686a4fd0393 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1=C5=A1=20Ml=C3=A1dek?= <t@mldk.cz>
Date: Wed, 31 Jan 2018 14:34:59 +0100
Subject: [PATCH] filetype detection

---
 requirements.in  |  3 ++-
 requirements.txt |  1 +
 robot.py         | 17 ++++++++++-------
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/requirements.in b/requirements.in
index b03f76c..dc174ff 100644
--- a/requirements.in
+++ b/requirements.in
@@ -1,3 +1,4 @@
 python-telegram-bot
 youtube-dl
-requests
\ No newline at end of file
+requests
+filetype
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index a851d07..80e0ac4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,6 +6,7 @@
 #
 certifi==2018.1.18        # via python-telegram-bot, requests
 chardet==3.0.4            # via requests
+filetype==1.0.0
 future==0.16.0            # via python-telegram-bot
 idna==2.6                 # via requests
 python-telegram-bot==9.0.0
diff --git a/robot.py b/robot.py
index 2c71994..dc8fe38 100755
--- a/robot.py
+++ b/robot.py
@@ -1,10 +1,12 @@
 #!/usr/bin/env python3
 import logging
 import os
+import re
 import shutil
 import sys
 from glob import glob
 
+import filetype
 import requests
 import youtube_dl
 from telegram import MessageEntity
@@ -50,18 +52,19 @@ def download_ydl(urls, date):
 
 
 def download_raw(url, date):
-    if 'api.telegram.org' in url \
-            and ('animation' in url or 'video' in url) \
-            and "mp4" not in url:
-        ext = '.mp4'
-    else:
-        ext = ''
-    local_filename = OUT_DIR + '/' + datestr(date) + '__' + url.split('/')[-1] + ext
+    local_filename = OUT_DIR + '/' + datestr(date) + '__' + url.split('/')[-1]
     r = requests.get(url, stream=True)
     with open(local_filename, 'wb') as f:
         for chunk in r.iter_content(chunk_size=1024):
             if chunk:
                 f.write(chunk)
+    if not re.match(r'.*\..{3,5}$', os.path.split(local_filename)[-1]):
+        kind = filetype.guess(local_filename)
+        if kind is None:
+            logger.error("File has no extension and could not be determined!")
+        else:
+            logger.info('Moving file without extension... %s?' % kind.extension)
+            shutil.move(local_filename, local_filename + '.' + kind.extension)
 
 
 # noinspection PyBroadException