From b623cccf3e19c987eb889c32f58b9df9abbcacf9 Mon Sep 17 00:00:00 2001
From: Eloi Manuel <eloi.manvie@gmail.com>
Date: Fri, 22 Dec 2023 12:21:15 +0100
Subject: [PATCH] Adding upgrade scripts

---
 script/upgradeManagedAccounts.s.sol    | 50 ++++++++++++++++++++++++++
 script/upgradeUpgradeableAccount.s.sol | 37 +++++++++++++++++++
 2 files changed, 87 insertions(+)
 create mode 100644 script/upgradeManagedAccounts.s.sol
 create mode 100644 script/upgradeUpgradeableAccount.s.sol

diff --git a/script/upgradeManagedAccounts.s.sol b/script/upgradeManagedAccounts.s.sol
new file mode 100644
index 0000000..22ec6ea
--- /dev/null
+++ b/script/upgradeManagedAccounts.s.sol
@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity =0.8.19;
+
+import {Script, console} from "forge-std/Script.sol";
+import {IEntryPoint} from "lib/account-abstraction/contracts/interfaces/IEntryPoint.sol";
+import {ManagedOpenfortAccount} from "../contracts/core/managed/ManagedOpenfortAccount.sol";
+import {ManagedOpenfortFactory} from "../contracts/core/managed/ManagedOpenfortFactory.sol";
+import {CheckOrDeployEntryPoint} from "script/aux/checkOrDeployEntryPoint.sol";
+
+contract ManagedOpenfortUpgrade is Script, CheckOrDeployEntryPoint {
+    uint256 internal deployPrivKey = vm.envUint("PK_PAYMASTER_OWNER_TESTNET");
+    address internal deployAddress = vm.addr(deployPrivKey);
+
+    address internal factoryAddress = 0x44A7d7B291834442EE3703bD8bB7f91eD6F2577E;
+    address internal oldImplementation = 0x36604309934A2Fc92C3445Cf4566b23b5b4BbAad;
+
+    function run()
+        public
+        returns (ManagedOpenfortAccount managedOpenfortAccountImpl, ManagedOpenfortFactory openfortFactory)
+    {
+        bytes32 versionSalt = vm.envBytes32("VERSION_SALT");
+        openfortFactory = ManagedOpenfortFactory(factoryAddress);
+
+        assert(openfortFactory.implementation() == oldImplementation);
+        address accountImpl = openfortFactory.implementation();
+        console.log("Old account implementation: ", accountImpl);
+
+        address exampleAccountAddress = openfortFactory.createAccountWithNonce(deployAddress, "1", true);
+        console.log("Example account address: ", exampleAccountAddress);
+
+        vm.startBroadcast(deployPrivKey);
+        // Create an acccount to serve as new implementation
+        managedOpenfortAccountImpl = new ManagedOpenfortAccount{salt: versionSalt}();
+
+        // Update the account implementation
+        openfortFactory.upgradeTo(address(managedOpenfortAccountImpl));
+        assert(openfortFactory.implementation() == address(managedOpenfortAccountImpl));
+
+        vm.stopBroadcast();
+
+        accountImpl = openfortFactory.implementation();
+        console.log("New account implementation: ", accountImpl);
+
+        // Create a managed account and get its address
+        address exampleAccountAddress2 = openfortFactory.createAccountWithNonce(deployAddress, "1", true);
+        console.log("Example account address 2: ", exampleAccountAddress2);
+
+        assert(exampleAccountAddress == exampleAccountAddress2);
+    }
+}
diff --git a/script/upgradeUpgradeableAccount.s.sol b/script/upgradeUpgradeableAccount.s.sol
new file mode 100644
index 0000000..8083652
--- /dev/null
+++ b/script/upgradeUpgradeableAccount.s.sol
@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity =0.8.19;
+
+import {Script, console} from "forge-std/Script.sol";
+import {IEntryPoint} from "lib/account-abstraction/contracts/interfaces/IEntryPoint.sol";
+import {UpgradeableOpenfortAccount} from "../contracts/core/upgradeable/UpgradeableOpenfortAccount.sol";
+import {UpgradeableOpenfortProxy} from "../contracts/core/upgradeable/UpgradeableOpenfortProxy.sol";
+import {CheckOrDeployEntryPoint} from "script/aux/checkOrDeployEntryPoint.sol";
+
+contract UpgradeableOpenfortDeploy is Script, CheckOrDeployEntryPoint {
+    uint256 internal deployPrivKey = vm.envUint("PK_PAYMASTER_OWNER_TESTNET");
+    address internal deployAddress = vm.addr(deployPrivKey);
+    IEntryPoint internal entryPoint;
+
+    address internal accountAddress = address(0);
+    address internal newImplementation = address(0);
+
+    event AccountImplementationDeployed(address indexed creator);
+
+    function run() public {
+        UpgradeableOpenfortProxy proxy = UpgradeableOpenfortProxy(payable(accountAddress));
+        UpgradeableOpenfortAccount account = UpgradeableOpenfortAccount(payable(accountAddress));
+
+        address accountImpl = proxy.implementation();
+        console.log("Old account implementation: ", accountImpl);
+
+        vm.startBroadcast(deployPrivKey);
+
+        // Update the account implementation
+        account.upgradeTo(newImplementation);
+
+        vm.stopBroadcast();
+
+        accountImpl = proxy.implementation();
+        console.log("New account implementation: ", accountImpl);
+    }
+}