From 5ceaf7f7fcc1f4b416fb9054e553897a1a8dddd4 Mon Sep 17 00:00:00 2001
From: Thrilok Kumar <thrilok2000@gmail.com>
Date: Wed, 29 Jul 2020 19:21:14 +0530
Subject: [PATCH] Added fee taker function

---
 contracts/connectors/1inch.sol | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/contracts/connectors/1inch.sol b/contracts/connectors/1inch.sol
index eeadcaf..767458a 100644
--- a/contracts/connectors/1inch.sol
+++ b/contracts/connectors/1inch.sol
@@ -5,6 +5,9 @@ import { TokenInterface , MemoryInterface, EventInterface} from "../common/inter
 import { Stores } from "../common/stores.sol";
 import { DSMath } from "../common/math.sol";
 
+import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/SafeERC20.sol";
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+
 interface OneInchInterace {
     function swap(
         TokenInterface fromToken,
@@ -62,6 +65,7 @@ interface OneProtoInterface {
 
 
 contract OneHelpers is Stores, DSMath {
+    using SafeERC20 for IERC20;
     /**
      * @dev Return  1Inch Address
      */
@@ -129,6 +133,30 @@ contract OneHelpers is Stores, DSMath {
         }
         return _tokens;
     }
+
+    function _transfer(address payable to, IERC20 token, uint _amt) internal {
+        address(token) == getEthAddr() ?
+            to.transfer(_amt) :
+            token.safeTransfer(to, _amt);
+    }
+
+    function takeFee(
+        address token,
+        uint amount,
+        address feeCollector,
+        uint feePercent
+    ) internal returns (uint leftAmt, uint feeAmount){
+        if (feeCollector != address(0)) {
+            feeAmount = wmul(amount, feePercent);
+            leftAmt = sub(amount, feeAmount);
+            uint feeCollectorAmt = wmul(feeAmount, 3 * 10 ** 17);
+            uint restAmt = sub(feeAmount, feeCollectorAmt);
+            _transfer(payable(feeCollector), IERC20(token), feeCollectorAmt);
+            _transfer(payable(getReferralAddr()), IERC20(token), restAmt); // TODO - change address
+        } else {
+            leftAmt = amount;
+        }
+    }
 }