From 225f9d92705af2b649782e3204b1b433a15d7114 Mon Sep 17 00:00:00 2001
From: Bob Mottram <bob@libreserver.org>
Date: Sun, 22 May 2022 20:55:11 +0100
Subject: [PATCH] Handle avatars and banners within HEAD

---
 daemon.py | 45 +++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 41 insertions(+), 4 deletions(-)

diff --git a/daemon.py b/daemon.py
index b534f1de0..439a9ae87 100644
--- a/daemon.py
+++ b/daemon.py
@@ -18598,13 +18598,47 @@ class PubServer(BaseHTTPRequestHandler):
         etag = None
         file_length = -1
 
-        if '/media/' in self.path:
+        if '/media/' in self.path or \
+           '/accounts/avatars/' in self.path or \
+           '/accounts/headers/' in self.path:
             if is_image_file(self.path) or \
                path_is_video(self.path) or \
                path_is_audio(self.path):
-                media_str = self.path.split('/media/')[1]
-                media_filename = \
-                    self.server.base_dir + '/media/' + media_str
+                if '/media/' in self.path:
+                    media_str = self.path.split('/media/')[1]
+                    media_filename = \
+                        self.server.base_dir + '/media/' + media_str
+                elif '/accounts/avatars/' in self.path:
+                    avatar_file = self.path.split('/accounts/avatars/')[1]
+                    if '/' not in avatar_file:
+                        self._404()
+                        return
+                    nickname = avatar_file.split('/')[0]
+                    avatar_file = avatar_file.split('/')[1]
+                    avatar_file_ext = avatar_file.split('.')[-1]
+                    # remove any numbers, eg. avatar123.png becomes avatar.png
+                    if avatar_file.startswith('avatar'):
+                        avatar_file = 'avatar.' + avatar_file_ext
+                    media_filename = \
+                        self.server.base_dir + '/accounts/' + \
+                        nickname + '@' + self.server.domain + '/' + \
+                        avatar_file
+                else:
+                    banner_file = self.path.split('/accounts/headers/')[1]
+                    if '/' not in banner_file:
+                        self._404()
+                        return
+                    nickname = banner_file.split('/')[0]
+                    banner_file = banner_file.split('/')[1]
+                    banner_file_ext = banner_file.split('.')[-1]
+                    # remove any numbers, eg. banner123.png becomes banner.png
+                    if banner_file.startswith('banner'):
+                        banner_file = 'banner.' + banner_file_ext
+                    media_filename = \
+                        self.server.base_dir + '/accounts/' + \
+                        nickname + '@' + self.server.domain + '/' + \
+                        banner_file
+
                 if os.path.isfile(media_filename):
                     check_path = media_filename
                     file_length = os.path.getsize(media_filename)
@@ -18632,6 +18666,9 @@ class PubServer(BaseHTTPRequestHandler):
                             except OSError:
                                 print('EX: do_HEAD unable to write ' +
                                       media_tag_filename)
+                else:
+                    self._404()
+                    return
 
         media_file_type = media_file_mime_type(check_path)
         self._set_headers_head(media_file_type, file_length,