From c8b8eef200c32abb66c4f49a27140b0d18393dfd Mon Sep 17 00:00:00 2001
From: m <mjkzyalt@gmail.com>
Date: Wed, 21 Dec 2022 15:39:59 -0600
Subject: [PATCH] add trusted ips (whitelist)

---
 src/client/component/party.cpp | 67 ++++++++++++++++++++++++++++++++--
 src/common/utils/io.cpp        | 21 +++++++++++
 src/common/utils/io.hpp        |  2 +
 3 files changed, 87 insertions(+), 3 deletions(-)

diff --git a/src/client/component/party.cpp b/src/client/component/party.cpp
index da9fb344..b3d6e79b 100644
--- a/src/client/component/party.cpp
+++ b/src/client/component/party.cpp
@@ -15,6 +15,7 @@
 
 #include "steam/steam.hpp"
 
+#include <utils/properties.hpp>
 #include <utils/string.hpp>
 #include <utils/info_string.hpp>
 #include <utils/cryptography.hpp>
@@ -309,6 +310,41 @@ namespace party
 
 		bool download_files(const game::netadr_s& target, const utils::info_string& info, bool allow_download);
 
+		std::string get_whitelist_json_path()
+		{
+			return (utils::properties::get_appdata_path() / "whitelist.json").generic_string();
+		}
+
+		bool get_whitelist_json_object(nlohmann::json* object)
+		{
+			std::string data;
+			if (!utils::io::read_file(get_whitelist_json_path(), &data))
+			{
+				return false;
+			}
+
+			try
+			{
+				*object = nlohmann::json::parse(data.data());
+			}
+			catch (const nlohmann::json::parse_error& ex)
+			{
+				menu_error(utils::string::va("%s\n", ex.what()));
+				return false;
+			}
+
+			return false;
+		}
+
+		std::string target_ip_to_string(const game::netadr_s& target)
+		{
+			return utils::string::va("%i.%i.%i.%i",
+				static_cast<int>(saved_info_response.host.ip[0]),
+				static_cast<int>(saved_info_response.host.ip[1]),
+				static_cast<int>(saved_info_response.host.ip[2]),
+				static_cast<int>(saved_info_response.host.ip[3]));
+		}
+
 		int user_download_response(game::hks::lua_State* state)
 		{
 			const auto response = state->m_apistack.base[0].v.boolean;
@@ -317,12 +353,36 @@ namespace party
 				return 0;
 			}
 
+			nlohmann::json obj;
+			if (!get_whitelist_json_object(&obj))
+			{
+				return false;
+			}
+
+			obj.insert(obj.end(), target_ip_to_string(saved_info_response.host));
+			utils::io::write_file_json(get_whitelist_json_path(), obj);
+
 			download_files(saved_info_response.host, saved_info_response.info_string, true);
 			return 1;
 		}
 
-		void confirm_user_download(const game::netadr_s& target, const utils::info_string& info)
+		bool should_user_confirm(const game::netadr_s& target, const utils::info_string& info)
 		{
+			nlohmann::json obj;
+			if (!get_whitelist_json_object(&obj))
+			{
+				return false;
+			}
+
+			const auto target_ip = target_ip_to_string(target);
+			for (const auto& [key, value] : obj.items())
+			{
+				if (value.is_string() && value.get<std::string>() == target_ip)
+				{
+					return false;
+				}
+			}
+
 			const auto LUI = ui_scripting::get_globals().get("LUI").as<ui_scripting::table>();
 			const auto yes_no_popup_func = LUI.get("yesnopopup").as<ui_scripting::function>();
 
@@ -334,6 +394,8 @@ namespace party
 			data_table.set("callback", user_download_response);
 
 			yes_no_popup_func(data_table);
+
+			return true;
 		}
 
 		bool needs_vid_restart = false;
@@ -350,9 +412,8 @@ namespace party
 
 				if (files.size() > 0)
 				{
-					if (!allow_download)
+					if (!allow_download && should_user_confirm(target, info))
 					{
-						confirm_user_download(target, info);
 						return true;
 					}
 
diff --git a/src/common/utils/io.cpp b/src/common/utils/io.cpp
index 5d898cc9..b05ee816 100644
--- a/src/common/utils/io.cpp
+++ b/src/common/utils/io.cpp
@@ -19,6 +19,27 @@ namespace utils::io
 		return std::ifstream(file).good();
 	}
 
+	bool write_file_json(const std::string& file, const nlohmann::json& object)
+	{
+		const auto pos = file.find_last_of("/\\");
+		if (pos != std::string::npos)
+		{
+			create_directory(file.substr(0, pos));
+		}
+
+		std::ofstream stream(
+			file, std::ios::binary | std::ofstream::out);
+
+		if (stream.is_open())
+		{
+			stream << object;
+			stream.close();
+			return true;
+		}
+
+		return false;
+	}
+
 	bool write_file(const std::string& file, const std::string& data, const bool append)
 	{
 		const auto pos = file.find_last_of("/\\");
diff --git a/src/common/utils/io.hpp b/src/common/utils/io.hpp
index 19e8c143..553930d8 100644
--- a/src/common/utils/io.hpp
+++ b/src/common/utils/io.hpp
@@ -3,12 +3,14 @@
 #include <string>
 #include <vector>
 #include <filesystem>
+#include <json.hpp>
 
 namespace utils::io
 {
 	bool remove_file(const std::string& file);
 	bool move_file(const std::string& src, const std::string& target);
 	bool file_exists(const std::string& file);
+	bool write_file_json(const std::string& file, const nlohmann::json& object);
 	bool write_file(const std::string& file, const std::string& data, bool append = false);
 	bool read_file(const std::string& file, std::string* data);
 	std::string read_file(const std::string& file);