ape-driven generic subscription management CLI

deployment/constants.py

@@ -12,12 +12,16 @@
 #
 # Domains
 #
+
 LYNX = "lynx"
 TAPIR = "tapir"
 MAINNET = "mainnet"
 
 SUPPORTED_TACO_DOMAINS = [LYNX, TAPIR, MAINNET]
 
+#
+# Nodes
+#
 
 LYNX_NODES = {
     # staking provider -> operator
@@ -39,3 +43,14 @@
 
 # Admin slot - https://eips.ethereum.org/EIPS/eip-1967#admin-address
 EIP1967_ADMIN_SLOT = 0xb53127684a568b3173ae13b9f8a6016e243e63b6e8ee1178d6a717850b5d6103
+
+
+#
+# Contracts
+#
+
+ACCESS_CONTROLLERS = [
+    "GlobalAllowList",
+    "OpenAccessAuthorizer",
+    "ManagedAllowList"
+]

deployment/options.py

@@ -0,0 +1,57 @@
+import click
+from eth_typing import ChecksumAddress
+
+from deployment.constants import (
+    ACCESS_CONTROLLERS,
+    SUPPORTED_TACO_DOMAINS
+)
+
+
+access_controller_option = click.option(
+    "--access-controller",
+    "-a",
+    help="global allow list or open access authorizer.",
+    type=click.Choice(ACCESS_CONTROLLERS),
+    required=True,
+)
+
+domain_option = click.option(
+    "--domain",
+    "-d",
+    help="TACo domain",
+    type=click.Choice(SUPPORTED_TACO_DOMAINS),
+    required=True,
+)
+
+ritual_id_option = click.option(
+    "--ritual-id",
+    "-r",
+    help="ID of the ritual",
+    required=True,
+    type=int
+)
+
+subscription_contract_option = click.option(
+    "--subscription-contract",
+    "-s",
+    help="Name of a subscription contract",
+    type=click.Choice(["BqETHSubscription"]),
+    required=True,
+)
+
+encryptor_slots_option = click.option(
+    "--encryptor-slots",
+    "-es",
+    help="Number of encryptor slots to pay for.",
+    required=True,
+    type=int
+)
+
+encryptors_option = click.option(
+    "--encryptors",
+    "-e",
+    help="List of encryptor addresses to remove.",
+    multiple=True,
+    required=True,
+    type=ChecksumAddress
+)

deployment/registry.py

@@ -5,13 +5,17 @@
 from pathlib import Path
 from typing import Dict, List, NamedTuple, Optional
 
+from ape import chain, project
 from ape.contracts import ContractInstance
-from ape import chain
 from eth_typing import ChecksumAddress
 from eth_utils import to_checksum_address
 from web3.types import ABI
 
-from deployment.utils import _load_json, get_contract_container
+from deployment.utils import (
+    _load_json,
+    get_contract_container,
+    registry_filepath_from_domain
+)
 
 ChainId = int
 ContractName = str
@@ -20,6 +24,10 @@
 STANDARD_REGISTRY_JSON_FORMAT = {"indent": 4, "separators": (",", ": ")}
 
 
+class NoContractFound(Exception):
+    """Raised when a contract is not found in the registry."""
+
+
 class RegistryEntry(NamedTuple):
     """Represents a single entry in a nucypher-style contract registry."""
 
@@ -295,3 +303,17 @@ def normalize_registry(filepath: Path):
     except Exception:
         print(f"Error when normalizing registry at {filepath}.")
         raise
+
+
+def get_contract(domain: str, contract_name: str) -> ContractInstance:
+    """Returns the contract instance for the contract name and domain."""
+    registry_filepath = registry_filepath_from_domain(domain=domain)
+    chain_id = project.chain_manager.chain_id
+    deployments = contracts_from_registry(filepath=registry_filepath, chain_id=chain_id)
+    try:
+        return deployments[contract_name]
+    except KeyError:
+        raise NoContractFound(
+            f"Contract '{contract_name}' not found in {domain} registry for chain {chain_id}. "
+            "Are you connected to the correct network + domain?"
+        )

deployment/utils.py

@@ -7,6 +7,7 @@
 from ape import networks, project
 from ape.contracts import ContractContainer, ContractInstance
 from ape_etherscan.utils import API_KEY_ENV_KEY_MAP
+
 from deployment.constants import ARTIFACTS_DIR
 from deployment.networks import is_local_network
 
@@ -136,8 +137,7 @@ def _get_dependency_contract_container(contract: str) -> ContractContainer:
             return contract_container
         except AttributeError:
             continue
-
-    raise ValueError(f"No contract found for {contract}")
+    raise ValueError(f"No contract found with name '{contract}'.")
 
 
 def get_contract_container(contract: str) -> ContractContainer:

scripts/manage_subscription.py

@@ -0,0 +1,174 @@
+import click
+from ape import Contract
+from ape.cli import account_option, ConnectedProviderCommand
+
+from deployment import registry
+from deployment.options import (
+    subscription_contract_option,
+    domain_option,
+    access_controller_option,
+    ritual_id_option,
+    encryptor_slots_option,
+    encryptors_option,
+)
+from deployment.params import Transactor
+from deployment.utils import check_plugins
+
+
+def _erc20_approve(
+        amount: int,
+        erc20: Contract,
+        receiver: Contract,
+        transactor: Transactor
+) -> None:
+    """Approve an ERC20 transfer."""
+    click.echo(
+        f"Approving transfer of {amount} {erc20.contract_type.name} "
+        f"to {receiver.contract_type.name}."
+    )
+    transactor.transact(
+        erc20.approve,
+        receiver.address,
+        amount
+    )
+
+
+def _calculate_slot_fees(
+        subscription_contract: Contract,
+        slots: int
+) -> int:
+    """Calculate the fees for a given number of encryptor slots."""
+    duration = subscription_contract.subscriptionPeriodDuration()
+    encryptor_fees = subscription_contract.encryptorFees(slots, duration)
+    total_fees = encryptor_fees
+    return total_fees
+
+
+@click.group()
+def cli():
+    """Subscription Management CLI"""
+
+
+@cli.command(cls=ConnectedProviderCommand)
+@account_option()
+@domain_option
+@subscription_contract_option
+@encryptor_slots_option
+@click.option(
+    "--period",
+    default=0,
+    help="Subscription billing period number to pay for.",
+)
+def pay_subscription(account, domain, subscription_contract, encryptor_slots, period):
+    """Pay for a new subscription period and initial encryptor slots."""
+    check_plugins()
+    transactor = Transactor(account=account)
+    subscription_contract = registry.get_contract(
+        contract_name=subscription_contract,
+        domain=domain
+    )
+    erc20 = Contract(subscription_contract.feeToken())
+    base_fees = subscription_contract.baseFees(period)
+    slot_fees = _calculate_slot_fees(
+        subscription_contract=subscription_contract,
+        slots=encryptor_slots
+    )
+    total_fees = base_fees + slot_fees
+    _erc20_approve(
+        amount=total_fees,
+        erc20=erc20,
+        receiver=subscription_contract,
+        transactor=transactor
+    )
+    click.echo(
+        f"Paying for subscription period #{period} "
+        f"with {encryptor_slots} encryptor slots."
+    )
+    transactor.transact(
+        subscription_contract.payForSubscription,
+        encryptor_slots
+    )
+
+
+@cli.command(cls=ConnectedProviderCommand)
+@account_option()
+@domain_option
+@subscription_contract_option
+@encryptor_slots_option
+def pay_slots(account, domain, subscription_contract, encryptor_slots):
+    """Pay for additional encryptor slots in the current billing period."""
+    check_plugins()
+    transactor = Transactor(account=account)
+    subscription_contract = registry.get_contract(
+        contract_name=subscription_contract,
+        domain=domain
+    )
+    erc20 = Contract(subscription_contract.feeToken())
+    fee = _calculate_slot_fees(
+        subscription_contract=subscription_contract,
+        slots=encryptor_slots
+    )
+    _erc20_approve(
+        amount=fee,
+        erc20=erc20,
+        receiver=subscription_contract,
+        transactor=transactor
+    )
+    click.echo(f"Paying for {encryptor_slots} new encryptor slots.")
+    transactor.transact(
+        subscription_contract.payForEncryptorSlots,
+        encryptor_slots
+    )
+
+
+@cli.command(cls=ConnectedProviderCommand)
+@account_option()
+@domain_option
+@ritual_id_option
+@access_controller_option
+@encryptors_option
+def add_encryptors(account, domain, ritual_id, access_controller, encryptors):
+    """Authorize encryptors to the access control contract for a ritual."""
+    access_controller = registry.get_contract(
+        contract_name=access_controller,
+        domain=domain
+    )
+    transactor = Transactor(account=account)
+    click.echo(
+        f"Adding {len(encryptors)} encryptors "
+        f"to the {access_controller} "
+        f"for ritual {ritual_id}."
+    )
+    transactor.transact(
+        access_controller.authorize,
+        ritual_id,
+        encryptors
+    )
+
+
+@cli.command(cls=ConnectedProviderCommand)
+@account_option()
+@domain_option
+@ritual_id_option
+@access_controller_option
+@encryptors_option
+def remove_encryptors(account, domain, ritual_id, access_controller, encryptors):
+    """Deauthorize encryptors from the access control contract for a ritual."""
+    transactor = Transactor(account=account)
+    access_controller = registry.get_contract(
+        contract_name=access_controller,
+        domain=domain
+    )
+    click.echo(
+        f"Removing {len(encryptors)} "
+        f"encryptors to the {access_controller} "
+        f"for ritual {ritual_id}."
+    )
+    transactor.transact(
+        access_controller.authorize,
+        ritual_id, encryptors
+    )
+
+
+if __name__ == "__main__":
+    cli()