From 1de20546182421e1280ec2b68c6d347abead7c54 Mon Sep 17 00:00:00 2001
From: Emilien <4016501+unixfox@users.noreply.github.com>
Date: Fri, 13 Dec 2024 20:08:57 +0100
Subject: [PATCH] add ability for invidious companion to check request from
 invidious

---
 config/config.example.yml                 |  1 +
 src/invidious/config.cr                   |  3 +++
 src/invidious/helpers/utils.cr            | 19 +++++++++++++++++++
 src/invidious/views/components/player.ecr |  9 ++++++---
 4 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/config/config.example.yml b/config/config.example.yml
index 2273c839..44b7c806 100644
--- a/config/config.example.yml
+++ b/config/config.example.yml
@@ -83,6 +83,7 @@ db:
 
 ##
 ## API key for Invidious companion
+## The size of the key needs to be more or equal to 16.
 ##
 ## Needed when invidious_companion is configured
 ##
diff --git a/src/invidious/config.cr b/src/invidious/config.cr
index 91622576..6508c847 100644
--- a/src/invidious/config.cr
+++ b/src/invidious/config.cr
@@ -250,6 +250,9 @@ class Config
       elsif config.invidious_companion_key == "CHANGE_ME!!"
         puts "Config: The value of 'invidious_companion_key' needs to be changed!!"
         exit(1)
+      elsif config.invidious_companion_key.size < 16
+        puts "Config: The value of 'invidious_companion_key' needs to be a size of 16 or more."
+        exit(1)
       end
     end
 
diff --git a/src/invidious/helpers/utils.cr b/src/invidious/helpers/utils.cr
index 4d9bb28d..85462eb8 100644
--- a/src/invidious/helpers/utils.cr
+++ b/src/invidious/helpers/utils.cr
@@ -383,3 +383,22 @@ def parse_link_endpoint(endpoint : JSON::Any, text : String, video_id : String)
   end
   return text
 end
+
+def encrypt_ecb_without_salt(data, key)
+  cipher = OpenSSL::Cipher.new("aes-128-ecb")
+  cipher.encrypt
+  cipher.key = key
+
+  io = IO::Memory.new
+  io.write(cipher.update(data))
+  io.write(cipher.final)
+  io.rewind
+
+  return io
+end
+
+def invidious_companion_encrypt(data)
+  timestamp = Time.utc.to_unix
+  encrypted_data = encrypt_ecb_without_salt("#{timestamp}|#{data}", CONFIG.invidious_companion_key)
+  return Base64.urlsafe_encode(encrypted_data)
+end
diff --git a/src/invidious/views/components/player.ecr b/src/invidious/views/components/player.ecr
index 975a25be..37f5f586 100644
--- a/src/invidious/views/components/player.ecr
+++ b/src/invidious/views/components/player.ecr
@@ -22,7 +22,8 @@
                audio_streams.each_with_index do |fmt, i|
                 src_url  = "/latest_version?id=#{video.id}&itag=#{fmt["itag"]}"
                 src_url += "&local=true" if params.local
-                src_url = video.invidious_companion["baseUrl"].as_s + src_url if (!CONFIG.invidious_companion.empty?)
+                src_url = video.invidious_companion["baseUrl"].as_s + src_url + 
+                            "&check=#{invidious_companion_encrypt(video.id)}" if (!CONFIG.invidious_companion.empty?)
 
                 bitrate = fmt["bitrate"]
                 mimetype = HTML.escape(fmt["mimeType"].as_s)
@@ -37,7 +38,8 @@
         <% else %>
             <% if params.quality == "dash"
                src_url = "/api/manifest/dash/id/" + video.id + "?local=true&unique_res=1"
-               src_url = video.invidious_companion["baseUrl"].as_s + src_url if (!CONFIG.invidious_companion.empty?)
+               src_url = video.invidious_companion["baseUrl"].as_s + src_url +
+                            "&check=#{invidious_companion_encrypt(video.id)}" if (!CONFIG.invidious_companion.empty?)
             %>
                 <source src="<%= src_url %>" type='application/dash+xml' label="dash">
             <% end %>
@@ -48,7 +50,8 @@
             fmt_stream.each_with_index do |fmt, i|
                 src_url  = "/latest_version?id=#{video.id}&itag=#{fmt["itag"]}"
                 src_url += "&local=true" if params.local
-                src_url = video.invidious_companion["baseUrl"].as_s + src_url if (!CONFIG.invidious_companion.empty?)
+                src_url = video.invidious_companion["baseUrl"].as_s + src_url +
+                            "&check=#{invidious_companion_encrypt(video.id)}" if (!CONFIG.invidious_companion.empty?)
 
                 quality = fmt["quality"]
                 mimetype = HTML.escape(fmt["mimeType"].as_s)