From 31411c22b6e0915e820f910c3a6a8ded507b5057 Mon Sep 17 00:00:00 2001
From: Derick Diaz <derickdiaz123@outlook.com>
Date: Sat, 20 Jan 2024 20:36:07 -0600
Subject: [PATCH] Added urlprotocol

---
 dnf5/commands/download/download.cpp | 33 ++++++++++++++++++++++++++++-
 dnf5/commands/download/download.hpp |  3 +++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/dnf5/commands/download/download.cpp b/dnf5/commands/download/download.cpp
index ced34949f8..e3edaae339 100644
--- a/dnf5/commands/download/download.cpp
+++ b/dnf5/commands/download/download.cpp
@@ -28,6 +28,8 @@ along with libdnf.  If not, see <https://www.gnu.org/licenses/>.
 #include <libdnf5/rpm/package_set.hpp>
 #include <libdnf5/utils/bgettext/bgettext-mark-domain.h>
 
+#include <algorithm>
+#include <format>
 #include <iostream>
 #include <map>
 
@@ -86,11 +88,25 @@ void DownloadCommand::set_argument_parser() {
     url->set_const_value("true");
     url->link_value(url_option);
 
+    auto urlprotocol = parser.add_new_named_arg("urlprotocol");
+    urlprotocol->set_long_name("urlprotocol");
+    urlprotocol->set_description("When running with --url, limit to specific protocols");
+    urlprotocol->set_parse_hook_func(
+        [&ctx](
+            [[maybe_unused]] ArgumentParser::NamedArg * arg, [[maybe_unused]] const char * option, const char * value) {
+            if (urlprotocol_valid_options.find(value) == urlprotocol_valid_options.end()) {
+                throw libdnf5::cli::ArgumentParserInvalidValueError(
+                    _M(std::format("Invalid urlprotocol option: {}", value)))
+            }
+            urlprotocol_option.emplace_back(value);
+        });
+
     cmd.register_named_arg(alldeps);
     create_destdir_option(*this);
     cmd.register_named_arg(resolve);
     cmd.register_positional_arg(keys);
     cmd.register_named_arg(url);
+    cmd.register_named_arg(urlprotocol);
 }
 
 void DownloadCommand::configure() {
@@ -161,10 +177,25 @@ void DownloadCommand::run() {
     }
 
     if (url_option->get_value()) {
+        // If no urlprotocols are specified, then any urlprotocol is acceptable
+        if (urlprotocol_option.empty()) {
+            urlprotocol_option = urlprotocol_valid_options
+        }
         for (auto & [nerva, pkg] : download_pkgs) {
             auto urls = pkg.get_remote_locations();
             libdnf_assert(!urls.empty(), "Failed to get mirror for package: \"{}\"", pkg.get_name());
-            std::cout << urls[0] << std::endl;
+            auto valid_url = std::find_if(urls.begin(), urls.end(), [&](std::string url) {
+                for (auto protocol : urlprotocols) {
+                    if (url.starts_with(protocol)) {
+                        return true;
+                    }
+                }
+                return false;
+            });
+            if (valid_url == urls.end()) {
+                libdnf_assert(true, "Failed to get mirror for package: \"{}\"", pkg.get_name());
+            }
+            std::cout << *valid_url << std::endl;
         }
         return;
     }
diff --git a/dnf5/commands/download/download.hpp b/dnf5/commands/download/download.hpp
index be9805f803..6aa7da2cd6 100644
--- a/dnf5/commands/download/download.hpp
+++ b/dnf5/commands/download/download.hpp
@@ -25,6 +25,7 @@ along with libdnf.  If not, see <https://www.gnu.org/licenses/>.
 #include <libdnf5/conf/option.hpp>
 
 #include <memory>
+#include <set>
 #include <vector>
 
 
@@ -40,6 +41,8 @@ class DownloadCommand : public Command {
     void run() override;
 
 private:
+    std::set<std::string> urlprotocol_valid_options = {"http", "https", "rsync", "ftp"};
+    std::set<std::string> urlprotocol_option = {};
     libdnf5::OptionBool * resolve_option{nullptr};
     libdnf5::OptionBool * alldeps_option{nullptr};
     libdnf5::OptionBool * url_option{nullptr};