From 2ed22cd3241010b1d8e4c289e4dd39680724275f Mon Sep 17 00:00:00 2001
From: jordaniza <j@jordaniza.com>
Date: Wed, 23 Oct 2024 11:06:58 +0400
Subject: [PATCH] Added the StakeInspector contract

---
 Makefile                          | 13 +++++++++++++
 script/DeployStakeInspector.s.sol | 21 +++++++++++++++++++++
 src/helpers/StakeInspector.sol    | 29 +++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+)
 create mode 100644 script/DeployStakeInspector.s.sol
 create mode 100644 src/helpers/StakeInspector.sol

diff --git a/Makefile b/Makefile
index a760692..5ea318f 100644
--- a/Makefile
+++ b/Makefile
@@ -39,6 +39,19 @@ ft-sepolia-fork :; forge test --match-contract TestE2EV2 \
 
 #### Deployments ####
 
+stakeinspector-preview-mode :; forge script DeployStakeInspector \
+  --rpc-url https://mainnet.mode.network \
+	-vvvvv
+
+stakeinspector-mode :; forge script DeployStakeInspector \
+  --rpc-url https://mainnet.mode.network \
+  --private-key $(DEPLOYMENT_PRIVATE_KEY) \
+  --broadcast \
+  --verify \
+  --verifier blockscout \
+  --verifier-url https://explorer.mode.network/api\? \
+  -vvvvv
+
 deploy-preview-mode-sepolia :; forge script script/Deploy.s.sol:Deploy \
   --rpc-url https://sepolia.mode.network \
 	--private-key $(DEPLOYMENT_PRIVATE_KEY) \
diff --git a/script/DeployStakeInspector.s.sol b/script/DeployStakeInspector.s.sol
new file mode 100644
index 0000000..66f45c9
--- /dev/null
+++ b/script/DeployStakeInspector.s.sol
@@ -0,0 +1,21 @@
+/// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.17;
+
+import {Script, console2 as console} from "forge-std/Script.sol";
+
+import {StakeInspector} from "src/helpers/StakeInspector.sol";
+
+contract DeployStakeInspector is Script {
+    address escrowMode = 0xff8AB822b8A853b01F9a9E9465321d6Fe77c9D2F;
+    address escrowBPT = 0x9c2eFe2a1FBfb601125Bb07a3D5bC6EC91F91e01;
+
+    function run() public {
+        vm.startBroadcast(vm.addr(vm.envUint("DEPLOYMENT_PRIVATE_KEY")));
+        StakeInspector stakeInspectorMode = new StakeInspector(escrowMode);
+        StakeInspector stakeInspectorBPT = new StakeInspector(escrowBPT);
+
+        vm.stopBroadcast();
+        console.log("StakeInspector for Mode deployed at: ", address(stakeInspectorMode));
+        console.log("StakeInspector for BPT deployed at: ", address(stakeInspectorBPT));
+    }
+}
diff --git a/src/helpers/StakeInspector.sol b/src/helpers/StakeInspector.sol
new file mode 100644
index 0000000..707fc4b
--- /dev/null
+++ b/src/helpers/StakeInspector.sol
@@ -0,0 +1,29 @@
+/// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.17;
+
+import {IVotingEscrowIncreasing} from "@escrow-interfaces/IVotingEscrowIncreasing.sol";
+
+interface IVotingEscrowIncreasingExtended is IVotingEscrowIncreasing {
+    function ownedTokens(address _account) external view returns (uint256[] memory tokenIds);
+}
+
+/// @notice Aggregate staked balances for an address given multiple veNFTs
+contract StakeInspector {
+    /// @notice The staking contract
+    IVotingEscrowIncreasingExtended public escrow;
+
+    constructor(address _escrow) {
+        escrow = IVotingEscrowIncreasingExtended(_escrow);
+    }
+
+    /// @notice Fetch the total underlying tokens staked across all the veNFTs owned by an account
+    function getTotalStaked(address _account) external view returns (uint256) {
+        uint256 totalStaked = 0;
+
+        uint256[] memory veNFTs = escrow.ownedTokens(_account);
+        for (uint256 i = 0; i < veNFTs.length; i++) {
+            totalStaked += escrow.locked(veNFTs[i]).amount;
+        }
+        return totalStaked;
+    }
+}