From 7b67a6d0877e181ee4ab118606a0c1ba6bbf5a79 Mon Sep 17 00:00:00 2001
From: Diavolo <edoardo.sanguineti222@gmail.com>
Date: Sat, 25 Jun 2022 20:22:13 +0200
Subject: [PATCH] DisableWeaponPickup & EnableWeaponPickup

---
 src/Components/Modules/AssetHandler.cpp |  2 +-
 src/Components/Modules/MapRotation.cpp  |  8 ++--
 src/Components/Modules/Weapon.cpp       | 58 +++++++++++++++++++++++--
 src/Components/Modules/Weapon.hpp       |  3 ++
 src/Game/Structs.hpp                    |  4 ++
 5 files changed, 67 insertions(+), 8 deletions(-)

diff --git a/src/Components/Modules/AssetHandler.cpp b/src/Components/Modules/AssetHandler.cpp
index a65259e2..1f2f4161 100644
--- a/src/Components/Modules/AssetHandler.cpp
+++ b/src/Components/Modules/AssetHandler.cpp
@@ -1,4 +1,4 @@
-#include <StdInclude.hpp>
+#include <STDInclude.hpp>
 
 #include "AssetInterfaces/IFont_s.hpp"
 #include "AssetInterfaces/IWeapon.hpp"
diff --git a/src/Components/Modules/MapRotation.cpp b/src/Components/Modules/MapRotation.cpp
index 30adf749..8c6658c4 100644
--- a/src/Components/Modules/MapRotation.cpp
+++ b/src/Components/Modules/MapRotation.cpp
@@ -226,11 +226,11 @@ namespace Components
 		// Ook, ook, eek
 		Logger::Warning(Game::CON_CHANNEL_SERVER, "You are using deprecated {}", (*SVMapRotationCurrent)->name);
 
-		RotationData currentRotation;
+		RotationData rotationCurrent;
 		try
 		{
 			Logger::Debug("Parsing {}", (*SVMapRotationCurrent)->name);
-			currentRotation.parse(data);
+			rotationCurrent.parse(data);
 		}
 		catch (const std::exception& ex)
 		{
@@ -239,14 +239,14 @@ namespace Components
 
 		Game::Dvar_SetString(*SVMapRotationCurrent, "");
 
-		if (currentRotation.getEntriesSize() == 0)
+		if (rotationCurrent.getEntriesSize() == 0)
 		{
 			Logger::Print(Game::CON_CHANNEL_SERVER, "{} is empty or contains invalid data. Restarting map\n", (*SVMapRotationCurrent)->name);
 			RestartCurrentMap();
 			return;
 		}
 
-		ApplyRotation(currentRotation);
+		ApplyRotation(rotationCurrent);
 	}
 
 	void MapRotation::RandomizeMapRotation()
diff --git a/src/Components/Modules/Weapon.cpp b/src/Components/Modules/Weapon.cpp
index 1b9719a3..82e2db4d 100644
--- a/src/Components/Modules/Weapon.cpp
+++ b/src/Components/Modules/Weapon.cpp
@@ -448,10 +448,57 @@ namespace Components
 		}
 	}
 
+	void __declspec(naked) Weapon::WeaponEntCanBeGrabbed_Stub()
+	{
+		using namespace Game;
+
+		__asm
+		{
+			cmp dword ptr [esp + 0x8], 0x0
+			jz touched
+
+			push 0x56E82C
+			retn
+
+		touched:
+			test dword ptr [edi + 0x2BC], PWF_DISABLE_WEAPON_PICKUP
+			jnz exit_func
+
+			// Game code
+			test eax, eax
+			jz continue_func
+
+		exit_func:
+			xor eax, eax
+			ret
+
+		continue_func:
+			push 0x56E84C
+			retn
+		}
+	}
+
+	void Weapon::AddScriptMethods()
+	{
+		Script::AddMethod("DisableWeaponPickup", [](Game::scr_entref_t entref)
+		{
+			const auto* ent = Game::GetPlayerEntity(entref);
+
+			ent->client->ps.weapCommon.weapFlags |= Game::PWF_DISABLE_WEAPON_PICKUP;
+		});
+
+		Script::AddMethod("EnableWeaponPickup", [](Game::scr_entref_t entref)
+		{
+			const auto* ent = Game::GetPlayerEntity(entref);
+
+			ent->client->ps.weapCommon.weapFlags &= ~Game::PWF_DISABLE_WEAPON_PICKUP;
+		});
+	}
+
 	Weapon::Weapon()
 	{
-		Weapon::PatchLimit();
-		Weapon::PatchConfigStrings();
+		PatchLimit();
+		PatchConfigStrings();
 
 		// Intercept weapon loading
 		AssetHandler::OnFind(Game::XAssetType::ASSET_TYPE_WEAPON, Weapon::WeaponFileLoad);
@@ -466,7 +513,7 @@ namespace Components
 
 		// Weapon swap fix
 		Utils::Hook::Nop(0x4B3670, 5);
-		Utils::Hook(0x57B4F0, LoadNoneWeaponHookStub).install()->quick();
+		Utils::Hook(0x57B4F0, LoadNoneWeaponHookStub, HOOK_JUMP).install()->quick();
 
 		// Don't load bounce sounds for now, it causes crashes
 		// TODO: Actually check the weaponfiles and/or reset the soundtable correctly!
@@ -477,5 +524,10 @@ namespace Components
 		// Clear weapons independently from fs_game
 		//Utils::Hook::Nop(0x452C1D, 2);
 		//Utils::Hook::Nop(0x452C24, 5);
+
+		AddScriptMethods();
+
+		AssertOffset(Game::playerState_s, Game::playerState_s::weapCommon.weapFlags, 0x2BC);
+		Utils::Hook(0x56E825, WeaponEntCanBeGrabbed_Stub, HOOK_JUMP).install()->quick();
 	}
 }
diff --git a/src/Components/Modules/Weapon.hpp b/src/Components/Modules/Weapon.hpp
index 3bd764f6..12dd99a2 100644
--- a/src/Components/Modules/Weapon.hpp
+++ b/src/Components/Modules/Weapon.hpp
@@ -25,5 +25,8 @@ namespace Components
 		static void ParseConfigStrings();
 		static int ParseWeaponConfigStrings();
 		static int ClearConfigStrings(void* dest, int value, int size);
+
+		static void WeaponEntCanBeGrabbed_Stub();
+		static void AddScriptMethods();
 	};
 }
diff --git a/src/Game/Structs.hpp b/src/Game/Structs.hpp
index 16a40cdd..06b420b2 100644
--- a/src/Game/Structs.hpp
+++ b/src/Game/Structs.hpp
@@ -1367,6 +1367,8 @@ namespace Game
 		PWF_DISABLE_WEAPON_SWAPPING = 1 << 11,
 		PWF_DISABLE_OFFHAND_WEAPONS = 1 << 12,
 		PWF_SWITCHING_TO_RIOTSHIELD = 1 << 13,
+		// IW5 flags backported
+		PWF_DISABLE_WEAPON_PICKUP = 1 << 16
 	};
 
 	struct playerState_s
@@ -1491,6 +1493,8 @@ namespace Game
 		int stunTime;
 	};
 
+	static_assert(sizeof(Game::playerState_s) == 0x311C);
+
 	enum LocSelInputState
 	{
 		LOC_SEL_INPUT_NONE = 0x0,