From db8b90487f1e5267419294d7a40fdb3c608caa5f Mon Sep 17 00:00:00 2001
From: "Gao, Ruiyuan" <905370712@qq.com>
Date: Wed, 25 Sep 2024 18:41:28 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=8A=A0=E5=AF=86=E9=9F=B3=E4=B9=90?=
 =?UTF-8?q?=E5=92=8C=E5=9B=BE=E7=89=87=E8=AE=BF=E9=97=AE=E9=93=BE=E6=8E=A5?=
 =?UTF-8?q?=20(#200)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* use basic auth, cannot work

* Revert "use basic auth, cannot work"

This reverts commit 16a9683855daed20efc4a9aef0628413380b62bc.

* use access key/code control

* Auto-format code 🧹🌟🤖

---------

Co-authored-by: Formatter [BOT] <runner@fv-az1766-921.lyuwioyq51hutffh0ei52p4blg.dx.internal.cloudapp.net>
---
 xiaomusic/httpserver.py | 43 +++++++++++++++++++++++++++++++++++++++--
 xiaomusic/xiaomusic.py  | 39 ++++++++++++++++++++++++++++++++++---
 2 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/xiaomusic/httpserver.py b/xiaomusic/httpserver.py
index b81fa52..9cf3aae 100644
--- a/xiaomusic/httpserver.py
+++ b/xiaomusic/httpserver.py
@@ -1,4 +1,5 @@
 import asyncio
+import hashlib
 import json
 import mimetypes
 import os
@@ -381,11 +382,46 @@ async def file_iterator(file_path, start, end):
             yield data
 
 
+def access_key_verification(file_path, key, code):
+    if config.disable_httpauth:
+        return True
+
+    log.debug(f"访问限制接收端[{file_path}, {key}, {code}]")
+    if key is not None:
+        current_key_bytes = key.encode("utf8")
+        correct_key_bytes = (
+            config.httpauth_username + config.httpauth_password
+        ).encode("utf8")
+        is_correct_key = secrets.compare_digest(correct_key_bytes, current_key_bytes)
+        if is_correct_key:
+            return True
+
+    if code is not None:
+        current_code_bytes = code.encode("utf8")
+        correct_code_bytes = (
+            hashlib.md5(
+                (
+                    file_path + config.httpauth_username + config.httpauth_password
+                ).encode("utf-8")
+            )
+            .hexdigest()
+            .encode("utf-8")
+        )
+        is_correct_code = secrets.compare_digest(correct_code_bytes, current_code_bytes)
+        if is_correct_code:
+            return True
+
+    return False
+
+
 range_pattern = re.compile(r"bytes=(\d+)-(\d*)")
 
 
 @app.get("/music/{file_path:path}")
-async def music_file(request: Request, file_path: str):
+async def music_file(request: Request, file_path: str, key: str = "", code: str = ""):
+    if not access_key_verification(request.url.path, key, code):
+        raise HTTPException(status_code=404, detail="File not found")
+
     absolute_path = os.path.abspath(config.music_path)
     absolute_file_path = os.path.normpath(os.path.join(absolute_path, file_path))
     if not absolute_file_path.startswith(absolute_path):
@@ -432,7 +468,10 @@ async def music_options():
 
 
 @app.get("/picture/{file_path:path}")
-async def get_picture(request: Request, file_path: str):
+async def get_picture(request: Request, file_path: str, key: str = "", code: str = ""):
+    if not access_key_verification(request.url.path, key, code):
+        raise HTTPException(status_code=404, detail="File not found")
+
     absolute_path = os.path.abspath(config.picture_cache_path)
     absolute_file_path = os.path.normpath(os.path.join(absolute_path, file_path))
     if not absolute_file_path.startswith(absolute_path):
diff --git a/xiaomusic/xiaomusic.py b/xiaomusic/xiaomusic.py
index ee162cc..4ea80bf 100644
--- a/xiaomusic/xiaomusic.py
+++ b/xiaomusic/xiaomusic.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 import asyncio
 import copy
+import hashlib
 import json
 import logging
 import math
@@ -107,6 +108,36 @@ class XiaoMusic:
         if self.config.conf_path == self.music_path:
             self.log.warning("配置文件目录和音乐目录建议设置为不同的目录")
 
+    def try_add_access_control_param(self, url):
+        if self.config.disable_httpauth:
+            return url
+
+        url_parts = urllib.parse.urlparse(url)
+        file_path = urllib.parse.unquote(url_parts.path)
+        correct_code = hashlib.md5(
+            (
+                file_path
+                + self.config.httpauth_username
+                + self.config.httpauth_password
+            ).encode("utf-8")
+        ).hexdigest()
+        self.log.debug(f"rewrite url: [{file_path}, {correct_code}]")
+
+        # make new url
+        parsed_get_args = dict(urllib.parse.parse_qsl(url_parts.query))
+        parsed_get_args.update({"code": correct_code})
+        encoded_get_args = urllib.parse.urlencode(parsed_get_args, doseq=True)
+        new_url = urllib.parse.ParseResult(
+            url_parts.scheme,
+            url_parts.netloc,
+            url_parts.path,
+            url_parts.params,
+            encoded_get_args,
+            url_parts.fragment,
+        ).geturl()
+
+        return new_url
+
     def init_config(self):
         self.music_path = self.config.music_path
         self.download_path = self.config.download_path
@@ -408,8 +439,8 @@ class XiaoMusic:
             if picture.startswith("/"):
                 picture = picture[1:]
             encoded_name = urllib.parse.quote(picture)
-            tags["picture"] = (
-                f"{self.hostname}:{self.public_port}/picture/{encoded_name}"
+            tags["picture"] = self.try_add_access_control_param(
+                f"{self.hostname}:{self.public_port}/picture/{encoded_name}",
             )
         return tags
 
@@ -451,7 +482,9 @@ class XiaoMusic:
         self.log.info(f"get_music_url local music. name:{name}, filename:{filename}")
 
         encoded_name = urllib.parse.quote(filename)
-        return f"{self.hostname}:{self.public_port}/music/{encoded_name}"
+        return self.try_add_access_control_param(
+            f"{self.hostname}:{self.public_port}/music/{encoded_name}",
+        )
 
     # 给前端调用
     def refresh_music_tag(self):