Add withdraw functions

.gitignore

@@ -6,6 +6,7 @@ out/
 !/broadcast
 /broadcast/*/31337/
 /broadcast/**/dry-run/
+lib/
 
 # Docs
 docs/

.gitmodules

@@ -0,0 +1,9 @@
+[submodule "lib/openzeppelin-contracts"]
+	path = lib/openzeppelin-contracts
+	url = https://github.com/openzeppelin/openzeppelin-contracts
+[submodule "lib/aave-v3-core"]
+	path = lib/aave-v3-core
+	url = https://github.com/aave/aave-v3-core
+[submodule "lib/forge-std"]
+	path = lib/forge-std
+	url = https://github.com/foundry-rs/forge-std

remappings.txt

@@ -1,4 +1,4 @@
 forge-std/=lib/forge-std/src/
 @chainlink/=/node_modules/@chainlink
-@openzeppelin/=/node_modules/@openzeppelin
-@aave/=/node_modules/@aave
+openzeppelin-contracts/=/home/trauslamen/Blockchain_Projects/gig-contracts/lib/openzeppelin-contracts
+aave/aave-v3-core/=/home/trauslamen/Blockchain_Projects/gig-contracts/lib/aave-v3-core

src/PaymentV1.sol

@@ -1,9 +1,9 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.19;
 
-import {IERC20} from "@chainlink/contracts-ccip/src/v0.8/vendor/openzeppelin-solidity/v4.8.0/contracts/token/ERC20/IERC20.sol";
-import {SafeERC20} from "@chainlink/contracts-ccip/src/v0.8/vendor/openzeppelin-solidity/v4.8.0/contracts/token/ERC20/utils/SafeERC20.sol";
-import {IPool} from "@aave/core-v3/contracts/interfaces/IPool.sol";
+import {IERC20} from "openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
+import {SafeERC20} from "openzeppelin-contracts/contracts/token/ERC20/utils/SafeERC20.sol";
+import {IPool} from "aave/aave-v3-core/contracts/interfaces/IPool.sol";
 
 contract USDCPayment {
     using SafeERC20 for IERC20;
@@ -13,21 +13,40 @@ contract USDCPayment {
     IPool public immutable i_pool;
     IERC20 public immutable i_usdc;
 
+    address public team_wallet; // safe wallet address
+
     error AmountIsZero();
+    error IncorrectAmount(uint256 amount);
     error TransferFailed();
     error IncorrectTokenAmount();
     error IncorrectAmountId();
     error NotOwner();
     error InvalidUsdcToken();
+    error IncorrectWalletAddress();
 
     event ProposalOpened(bytes32 id, uint256 amount, address initiator);
     event ProposalClosed(bytes32 id, address freelancer);
+    event WalletChanged(address curr_wallet, address new_wallet);
+    event WithdrawUSDC(address receiver, uint256 amount);
+    event WithdrawETH(address receiver, uint256 amount);
+
+    modifier onlyTeam() {
+        if (msg.sender != team_wallet) {
+            revert IncorrectWalletAddress();
+        }
+        _;
+    }
 
-    constructor(address _addressPool, address _usdcAddress) {
+    constructor(
+        address _addressPool,
+        address _usdcAddress,
+        address _teamAddress
+    ) {
         if (_usdcAddress == address(0)) revert InvalidUsdcToken();
         i_pool = IPool(_addressPool);
         i_usdc = IERC20(_usdcAddress);
-        i_usdc.safeApprove(_addressPool, type(uint256).max);
+        i_usdc.safeIncreaseAllowance(_addressPool, type(uint256).max);
+        team_wallet = _teamAddress;
     }
 
     function openProposal(uint256 _amount) external {
@@ -65,6 +84,33 @@ contract USDCPayment {
         i_pool.withdraw(address(i_usdc), proposalAmount, receiver);
     }
 
+    function withdrawUSDC(uint256 _amount) public onlyTeam {
+        (uint256 totalCollateralBase, , , , , ) = i_pool.getUserAccountData(
+            address(this)
+        ); // get all lended money from AAVE Pool
+        if (_amount > totalCollateralBase) {
+            revert IncorrectAmount(_amount);
+        }
+        i_usdc.safeTransfer(msg.sender, _amount);
+
+        emit WithdrawUSDC(msg.sender, _amount);
+    }
+
+    function withdrawETH() public onlyTeam {
+        uint256 balance = address(this).balance;
+        (bool success, ) = msg.sender.call{value: balance}("");
+        if (!success) {
+            revert TransferFailed();
+        }
+
+        emit WithdrawETH(msg.sender, balance);
+    }
+
+    function changeTeamWallet(address _newTeamWallet) public onlyTeam {
+        team_wallet = _newTeamWallet;
+        emit WalletChanged(team_wallet, _newTeamWallet);
+    }
+
     function getBalance(address _tokenAddress) external view returns (uint256) {
         return IERC20(_tokenAddress).balanceOf(address(this));
     }

test/PaymentTest.t.sol

@@ -4,18 +4,20 @@ pragma solidity ^0.8.19;
 import {Test} from "forge-std/Test.sol";
 import {USDCPayment} from "../src/PaymentV1.sol";
 import {IUSDC} from "../src/interfaces/IUSDC.sol";
-import {IERC20} from "@chainlink/contracts-ccip/src/v0.8/vendor/openzeppelin-solidity/v4.8.0/contracts/token/ERC20/IERC20.sol";
-import {IPool} from "@aave/core-v3/contracts/interfaces/IPool.sol";
+import {IERC20} from "openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
+import {IPool} from "aave/aave-v3-core/contracts/interfaces/IPool.sol";
 
 contract USDCPaymentTest is Test {
     USDCPayment private usdcPayment;
     address public constant USDC = 0xaf88d065e77c8cC2239327C5EDb3A432268e5831; // Arbitrum mainnet USDC token
     address public constant POOL = 0x794a61358D6845594F94dc1DB02A252b5b4814aD; // Arbitrum mainnet Pool
     uint256 public constant TESTING_AMOUNT = 1e6;
 
+    address public testWallet = makeAddr("TEST_WALLET");
+
     function setUp() public {
         vm.createSelectFork(vm.envString("ARBITRUM_RPC_URL"));
-        usdcPayment = new USDCPayment(POOL, USDC);
+        usdcPayment = new USDCPayment(POOL, USDC, testWallet);
     }
 
     function mintTokens() public {
@@ -224,4 +226,13 @@ contract USDCPaymentTest is Test {
             "Need to delete proposal budget after closing."
         );
     }
+
+    function test_RevertWithdrawIfNotOwner() public {
+        address tester = makeAddr("test_teamWallet");
+        vm.startPrank(tester);
+        vm.expectRevert(USDCPayment.IncorrectWalletAddress.selector);
+        usdcPayment.withdrawUSDC(TESTING_AMOUNT);
+        vm.expectRevert(USDCPayment.IncorrectWalletAddress.selector);
+        usdcPayment.withdrawETH();
+    }
 }