Merge pull request #25 from infosec-us-team/main

Feature: Adding powerful, granular, and fully customizable logging capabilities to PoCs.
immunefi-team · Nov 16, 2023 · 3c0e232 · 3c0e232
2 parents 6fdc278 + def6f76
commit 3c0e232
Show file tree

Hide file tree

Showing 6 changed files with 1,903 additions and 15 deletions.
diff --git a/README.md b/README.md
@@ -85,7 +85,7 @@ forge test -vv --match-path test/[test_name]
 
 |       Categorisation       | Branch | Source | Documentation |
 | -------------------------- | -------- | ------ | ------------- |
-| Default                 | [default](https://github.com/immunefi-team/forge-poc-templates/tree/default)        | | |
+| Default                    | [default](https://github.com/immunefi-team/forge-poc-templates/tree/default)        | | |
 | Reentrancy                 | [reentrancy](https://github.com/immunefi-team/forge-poc-templates/tree/reentrancy)        | [Source](./src/reentrancy/Reentrancy.sol)                         | [Readme](./src/reentrancy/README.md)              |
 | Flash Loan                 | [flash_loan](https://github.com/immunefi-team/forge-poc-templates/tree/flash_loan)         | [Source](./src/flashloan/FlashLoan.sol)                         | [Readme](./src/flashloan/README.md)               |
 | Price Manipulation         | [price_manipulation](https://github.com/immunefi-team/forge-poc-templates/tree/price_manipulation) | [Source](./src/pricemanipulation/PriceManipulation.sol)       | [Readme](./src/pricemanipulation/README.md)       |
@@ -117,6 +117,7 @@ We sincerely appreciate contributions to Immunefi's templates. Please take the t
 * [@realgmhacker](https://twitter.com/realgmhacker)
 * [@_iphelix](https://twitter.com/_iphelix)
 * [0xlead](https://github.com/0xlead)
+* [@infosec_us_team](https://twitter.com/infosec_us_team)
 
 
 ## All set!

diff --git a/src/PoC.sol b/src/PoC.sol
@@ -4,13 +4,14 @@ pragma solidity ^0.8.13;
 import "forge-std/Test.sol";
 
 import "./tokens/Tokens.sol";
+import "./log/Log.sol";
 
 struct TokenBalance {
     IERC20 token;
     int256 amount;
 }
 
-contract PoC is Test, Tokens {
+contract PoC is Test, Tokens, Log {
     // For snapshotting users' balances
     mapping(address => TokenBalance[][]) public tokensBalance;
     // For resolving addresses to aliases
@@ -80,10 +81,10 @@ contract PoC is Test, Tokens {
     function printBalance(address _user, uint256 _index) public view {
         string memory resolvedAddress = _resolveAddress(_user);
         if (logLevel == 1) {
-            console.log("~~~ Balance of [%s] at block #%s", resolvedAddress, block.number);
-            console.log("-----------------------------------------------------------------------------------------");
-            console.log("             Token address                    |       Symbol  |       Balance");
-            console.log("-----------------------------------------------------------------------------------------");
+            _log("~~~ Balance of [%s] at block #%s", resolvedAddress, block.number);
+            _log("-----------------------------------------------------------------------------------------");
+            _log("             Token address                    |       Symbol  |       Balance");
+            _log("-----------------------------------------------------------------------------------------");
         }
         for (uint256 j = 0; j < tokensBalance[_user][_index].length; j++) {
             uint256 balance = uint256(tokensBalance[_user][_index][j].amount);
@@ -104,13 +105,13 @@ contract PoC is Test, Tokens {
             string memory template;
             if (logLevel == 1) {
                 template = string.concat("%s\t|\t", symbol, "\t|\t%s.%s");
-                console.log(template, address(tokensBalance[_user][_index][j].token), integer_part, fractional_part);
+                _log(template, address(tokensBalance[_user][_index][j].token), integer_part, fractional_part);
             } else if (logLevel == 0) {
                 template = string.concat("--- ", symbol, " balance of [%s]:\t%s.%s", " ---");
-                console.log(template, resolvedAddress, integer_part, fractional_part);
+                _log(template, resolvedAddress, integer_part, fractional_part);
             }
         }
-        console.log();
+        _log();
     }
 
     /**
@@ -119,10 +120,10 @@ contract PoC is Test, Tokens {
      */
     function printProfit(address _user) public view {
         string memory resolvedAddress = _resolveAddress(_user);
-        console.log("~~~ Profit for [%s]", resolvedAddress);
-        console.log("-----------------------------------------------------------------------------------------");
-        console.log("             Token address                    |       Symbol  |       Profit");
-        console.log("-----------------------------------------------------------------------------------------");
+        _log("~~~ Profit for [%s]", resolvedAddress);
+        _log("-----------------------------------------------------------------------------------------");
+        _log("             Token address                    |       Symbol  |       Profit");
+        _log("-----------------------------------------------------------------------------------------");
         for (uint256 j = 0; j < tokensBalance[_user][0].length; j++) {
             int256 profit = tokensBalance[_user][tokensBalance[_user].length - 1][j].amount
                 - int256(tokensBalance[_user][0][j].amount);
@@ -146,9 +147,9 @@ contract PoC is Test, Tokens {
             // Generate template string
             string memory template = string.concat("%s\t|\t", symbol, "\t|\t", sign, "%s.%s");
 
-            console.log(template, address(tokensBalance[_user][0][j].token), integer_part, fractional_part);
+            _log(template, address(tokensBalance[_user][0][j].token), integer_part, fractional_part);
         }
-        console.log();
+        _log();
     }
 
     /**