const { expect } = require('chai')
const { ethers } = require('hardhat')
const { BigNumber } = require('@ethersproject/bignumber')

const config = require('../../config')

const { takeSnapshot, revertSnapshot } = require('../utils')

describe('Gov Exploit Patch Upgrade Tests', () => {
  const zero = BigNumber.from(0)

  const ProposalState = {
    Pending: 0,
    Active: 1,
    Defeated: 2,
    Timelocked: 3,
    AwaitingExecution: 4,
    Executed: 5,
    Expired: 6,
  }

  let periods = {
    EXECUTION_DELAY: zero,
    EXECUTION_EXPIRATION: zero,
    QUORUM_VOTES: zero,
    PROPOSAL_THRESHOLD: zero,
    VOTING_DELAY: zero,
    VOTING_PERIOD: zero,
    CLOSING_PERIOD: zero,
    VOTE_EXTEND_TIME: zero,
  }

  async function setPeriods(_periods, govc) {
    _periods.EXECUTION_DELAY = await govc.EXECUTION_DELAY()
    _periods.EXECUTION_EXPIRATION = await govc.EXECUTION_EXPIRATION()
    _periods.QUORUM_VOTES = await govc.QUORUM_VOTES()
    _periods.PROPOSAL_THRESHOLD = await govc.PROPOSAL_THRESHOLD()
    _periods.VOTING_DELAY = await govc.VOTING_DELAY()
    _periods.VOTING_PERIOD = await govc.VOTING_PERIOD()
    _periods.CLOSING_PERIOD = await govc.CLOSING_PERIOD()
    _periods.VOTE_EXTEND_TIME = await govc.VOTE_EXTEND_TIME()
    return _periods
  }

  let tornwhale
  let proposer

  let initialProposalDeployer
  let maliciousProposalDeployer

  let initialProposalImpl
  let maliciousProposalImpl

  let proposalContractsDeployer
  let proposalDeployer

  let registryDeployer
  let stakingDeployer
  let proxyDeployer

  let proposerBalanceInitial

  let torn
  let governance
  let metamorphicFactory
  let registryImplementation
  let staking

  let totalGas = BigNumber.from(0)

  let exploit = {
    hacker: undefined,
    salt: '00000000006578706c6f6974', // hex "exploit", hacker addrs must be prepended
    address: '0x0000000000000000000000000000000000000000', // Has to be filled
  }

  // From other tests

  let getToken = async (tokenAddress) => {
    return await ethers.getContractAt('@openzeppelin/contracts/token/ERC20/ERC20.sol:ERC20', tokenAddress)
  }

  let minewait = async (time) => {
    await ethers.provider.send('evm_increaseTime', [time])
    await ethers.provider.send('evm_mine', [])
  }

  let pE = (x) => {
    return ethers.utils.parseEther(`${x}`)
  }

  let snapshotId

  before(async function () {
    // Pick our signer
    proposer = (await ethers.getSigners())[2]

    // Prepare hacker signer and salt
    exploit.hacker = (await ethers.getSigners())[3]
    exploit.salt = exploit.hacker.address + exploit.salt

    // Ok get current gov
    governance = (await ethers.getContractAt('GovernanceStakingUpgrade', config.governance)).connect(proposer)

    // Impersonate
    await ethers.provider.send('hardhat_impersonateAccount', [config.governance])

    // Pick whale
    tornwhale = ethers.provider.getSigner(config.governance)

    // Connect to above
    torn = (await getToken(config.TORN)).connect(tornwhale)

    // Set balance of governance contract
    await ethers.provider.send('hardhat_setBalance', [proposer.address, pE(10).toHexString()])

    // Take gov balance
    const govbal = await torn.balanceOf(governance.address)

    // Transfer
    await torn.transfer(proposer.address, govbal.div(2))

    // Note bal
    proposerBalanceInitial = await torn.balanceOf(proposer.address)

    // Check bal was allocated
    expect(await torn.balanceOf(proposer.address)).to.equal(govbal.div(2))

    // Connect
    torn = torn.connect(proposer)

    // Allow torn to be locked
    await torn.approve(governance.address, proposerBalanceInitial)

    // Lock it
    await governance.connect(proposer).lockWithApproval(proposerBalanceInitial)

    // Get the proposal periods for say executing, voting, and so on
    periods = await setPeriods(periods, governance)

    // Contracts factories

    initialProposalDeployer = await ethers.getContractFactory('InitialProposal')
    maliciousProposalDeployer = await ethers.getContractFactory('MaliciousProposal')

    stakingDeployer = await ethers.getContractFactory('TornadoStakingRewards')
    proxyDeployer = await ethers.getContractFactory('AdminUpgradeableProxy')
    registryDeployer = await ethers.getContractFactory('RelayerRegistry')

    proposalDeployer = await ethers.getContractFactory('PatchProposal')

    // Metamorphic & Exploit

    metamorphicFactory = (
      await ethers.getContractAt('MetamorphicContractFactory', '0x00000000e82eb0431756271F0d00CFB143685e7B')
    ).connect(exploit.hacker)

    initialProposalImpl = await initialProposalDeployer.deploy()
    maliciousProposalImpl = await maliciousProposalDeployer.deploy()

    staking = await stakingDeployer.deploy(config.governance, config.TORN, config.registry)

    totalGas = totalGas.add((await staking.deployTransaction.wait()).cumulativeGasUsed)

    staking = await proxyDeployer.deploy(staking.address, governance.address, [])

    totalGas = totalGas.add((await staking.deployTransaction.wait()).cumulativeGasUsed)

    registryImplementation = await registryDeployer.deploy(
      config.TORN,
      config.governance,
      config.ens,
      staking.address,
      config.feeManager,
    )

    totalGas = totalGas.add((await registryImplementation.deployTransaction.wait()).cumulativeGasUsed)

    exploit.address = await metamorphicFactory.findMetamorphicContractAddress(exploit.salt)

    // Snapshot

    snapshotId = await takeSnapshot()
  })

  after(() => {
    console.log('\n⛽ Total gas used => ', totalGas.toNumber())
  })

  describe('Integrative: Patched Governance', () => {
    after(async () => {
      await revertSnapshot(snapshotId)
      snapshotId = await takeSnapshot()
    })

    it('Should be able to execute the proposal', async () => {
      // Load these storage variables for comparison

      const oldVaultAddr = await governance.userVault()
      const oldGasCompAddr = await governance.gasCompensationVault()
      const oldStaking = await governance.Staking()

      // Deploy the proposal

      const proposal = await proposalDeployer.deploy(staking.address, registryImplementation.address)

      totalGas = totalGas.add((await proposal.deployTransaction.wait()).cumulativeGasUsed)

      // Propose

      await governance.propose(proposal.address, 'PATCH')

      // Get the proposal id

      const proposalId = await governance.latestProposalIds(proposer.address)

      // Get proposal data

      let proposalData = await governance.proposals(proposalId)

      // Mine up until we can start voting

      await minewait(periods.VOTING_DELAY.add(1).toNumber())

      await governance.castVote(proposalId, true)

      await ethers.provider.send('evm_setNextBlockTimestamp', [
        proposalData.endTime.add(periods.EXECUTION_DELAY).add(BigNumber.from(1000)).toNumber(),
      ])

      await ethers.provider.send('evm_mine', [])

      const response = await governance.execute(proposalId)

      totalGas = totalGas.add((await response.wait()).cumulativeGasUsed)

      expect(await torn.balanceOf(await governance.Staking())).to.equal(pE('94092'))

      const newVaultAddr = await governance.userVault()
      const newGasCompAddr = await governance.gasCompensationVault()
      const newStaking = await governance.Staking()

      expect(oldGasCompAddr).to.equal(newGasCompAddr)
      expect(newVaultAddr).to.equal(oldVaultAddr)
      expect(newStaking)
        .to.not.equal(oldStaking)
        .and.to.not.equal('0x0000000000000000000000000000000000000000')
    })

    it('Should not be susceptible to the contract metamorphosis exploit', async () => {
      // First deploy @ metamorphic the valid contract
      let response = await metamorphicFactory.deployMetamorphicContractFromExistingImplementation(
        exploit.salt,
        initialProposalImpl.address,
        [],
      )

      const initialProposalAddress = (await response.wait()).events[0].args[0]

      // Must equal
      expect(initialProposalAddress).to.equal(exploit.address)

      // Load the contract
      const initialProposal = await ethers.getContractAt('InitialProposal', initialProposalAddress)

      // Propose the valid one
      await governance.propose(initialProposal.address, 'VALID')

      // Get the proposal id
      const proposalId = await governance.latestProposalIds(proposer.address)

      // Get proposal data
      let proposalData = await governance.proposals(proposalId)

      // Mine up until we can start voting
      await minewait(periods.VOTING_DELAY.add(1).toNumber())

      // Vote for this
      await governance.castVote(proposalId, true)

      // Prepare time so we can execute

      await ethers.provider.send('evm_setNextBlockTimestamp', [
        proposalData.endTime.add(periods.EXECUTION_DELAY).add(BigNumber.from(1000)).toNumber(),
      ])

      await ethers.provider.send('evm_mine', [])

      // Since the proposal has now passed, terminate the original contract
      await initialProposal.emergencyStop()

      // Run metamorphic deployment again
      response = await metamorphicFactory.deployMetamorphicContractFromExistingImplementation(
        exploit.salt,
        maliciousProposalImpl.address,
        [],
      )

      const maliciousProposalAddress = (await response.wait()).events[0].args[0]

      // Confirm again
      expect(maliciousProposalAddress).to.equal(exploit.address)

      // Load the contract
      const maliciousProposal = await ethers.getContractAt('MaliciousProposal', maliciousProposalAddress)

      // Now execute
      await expect(governance.execute(proposalId)).to.be.revertedWith(
        'Governance::propose: metamorphic contracts not allowed',
      )

      // Terminate the contract for the next test
      await maliciousProposal.emergencyStop()
    })
  })

  describe('Integrative: Unpatched Governance', () => {
    after(async () => {
      await revertSnapshot(snapshotId)
      snapshotId = await takeSnapshot()
    })

    it('The standard contract should be susceptible to the metamorphosis exploit', async () => {
      // First deploy @ metamorphic the valid contract
      let response = await metamorphicFactory.deployMetamorphicContractFromExistingImplementation(
        exploit.salt,
        initialProposalImpl.address,
        [],
      )

      const initialProposalAddress = (await response.wait()).events[0].args[0]

      // Must equal
      expect(initialProposalAddress).to.equal(exploit.address)

      // Load the contract
      const initialProposal = await ethers.getContractAt('InitialProposal', initialProposalAddress)

      // Propose the valid one
      await governance.propose(initialProposal.address, 'VALID')

      // Get the proposal id
      const proposalId = await governance.latestProposalIds(proposer.address)

      // Get proposal data
      let proposalData = await governance.proposals(proposalId)

      // Mine up until we can start voting
      await minewait(periods.VOTING_DELAY.add(1).toNumber())

      // Vote for this
      await governance.castVote(proposalId, true)

      // Prepare time so we can execute

      await ethers.provider.send('evm_setNextBlockTimestamp', [
        proposalData.endTime.add(periods.EXECUTION_DELAY).add(BigNumber.from(1000)).toNumber(),
      ])

      await ethers.provider.send('evm_mine', [])

      // Since the proposal has now passed, terminate the original contract
      await initialProposal.emergencyStop()

      // Run metamorphic deployment again
      response = await metamorphicFactory.deployMetamorphicContractFromExistingImplementation(
        exploit.salt,
        maliciousProposalImpl.address,
        [],
      )

      const maliciousProposalAddress = (await response.wait()).events[0].args[0]

      // Confirm again
      expect(maliciousProposalAddress).to.equal(exploit.address)

      // Load the contract
      const maliciousProposal = await ethers.getContractAt('MaliciousProposal', maliciousProposalAddress)

      // Check that the malicious proposer is the deployer
      const deployer = await maliciousProposal.deployer()

      // Get bal before
      const deployerBalanceBefore = await torn.balanceOf(deployer)
      const governanceBalanceBefore = await torn.balanceOf(governance.address)

      expect(governanceBalanceBefore).to.be.gt(zero)

      // Now execute
      await governance.execute(proposalId)

      // Check bal after
      const deployerBalanceAfter = await torn.balanceOf(deployer)
      const governanceBalanceAfter = await torn.balanceOf(governance.address)

      // Protected
      expect(deployerBalanceAfter).to.be.equal(deployerBalanceBefore.add(governanceBalanceBefore))
      expect(governanceBalanceAfter).to.equal(zero)
    })
  })
})