From dfc5b987b247b334566ac27ba8fc0caa74926048 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20N=C3=BA=C3=B1ez?= <david@nucypher.com>
Date: Tue, 12 Nov 2019 02:32:50 +0100
Subject: [PATCH] First pass at `nucypher stake preallocations
 --withdraw-tokens` and `--withdraw-eth`

---
 nucypher/blockchain/eth/actors.py | 24 +++++++++++++++++-
 nucypher/cli/characters/stake.py  | 42 ++++++++++++++++++++++++++++++-
 2 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/nucypher/blockchain/eth/actors.py b/nucypher/blockchain/eth/actors.py
index 3b9b4bed6..714e9dd37 100644
--- a/nucypher/blockchain/eth/actors.py
+++ b/nucypher/blockchain/eth/actors.py
@@ -894,7 +894,7 @@ class Staker(NucypherTokenActor):
     @only_me
     @save_receipt
     def withdraw(self, amount: NU) -> dict:
-        """Withdraw tokens (assuming they're unlocked)"""
+        """Withdraw tokens from StakingEscrow (assuming they're unlocked)"""
         if self.is_contract:
             receipt = self.preallocation_escrow_agent.withdraw_as_staker(value=int(amount))
         else:
@@ -902,6 +902,28 @@ class Staker(NucypherTokenActor):
                                                   amount=int(amount))
         return receipt
 
+    @only_me
+    @save_receipt
+    def withdraw_preallocation_tokens(self, amount: NU) -> dict:
+        """Withdraw tokens from PreallocationEscrow (assuming they're unlocked)"""
+        if amount <= 0:
+            raise ValueError(f"Don't try to withdraw {amount}.")
+        if self.is_contract:
+            receipt = self.preallocation_escrow_agent.withdraw_tokens(value=int(amount))
+        else:
+            raise TypeError("This method can only be used when staking via a contract")
+        return receipt
+
+    @only_me
+    @save_receipt
+    def withdraw_preallocation_eth(self) -> dict:
+        """Withdraw ETH from PreallocationEscrow"""
+        if self.is_contract:
+            receipt = self.preallocation_escrow_agent.withdraw_eth()
+        else:
+            raise TypeError("This method can only be used when staking via a contract")
+        return receipt
+
 
 class Worker(NucypherTokenActor):
     """
diff --git a/nucypher/cli/characters/stake.py b/nucypher/cli/characters/stake.py
index a34ad9981..30db02928 100644
--- a/nucypher/cli/characters/stake.py
+++ b/nucypher/cli/characters/stake.py
@@ -643,6 +643,9 @@ def collect_reward(click_config,
 @stake.command('preallocation')
 @_stake_options
 @click.option('--status', help="View balance and lock information of preallocation contract", is_flag=True)
+@click.option('--withdraw-tokens', help="Withdraw tokens to beneficiary address", is_flag=True)
+@click.option('--withdraw-eth', help="Withdraw ETH to beneficiary address", is_flag=True)
+@click.option('--force', help="Don't ask for confirmation", is_flag=True)
 @nucypher_click_config
 def preallocation(click_config,
 
@@ -651,11 +654,18 @@ def preallocation(click_config,
                   beneficiary_address, allocation_filepath,
 
                   # Preallocation options
-                  status):
+                  status, withdraw_tokens, withdraw_eth,
+
+                  # Other
+                  force):
     """
     Claim rewards and fees collected by a preallocation contract.
     """
 
+    flags = (status, withdraw_tokens, withdraw_eth)
+    if sum(flags) != 1:  # i.e., only one flag must be active
+        raise click.BadArgumentUsage(f"Either --status, --withdraw-tokens, or --withdraw-eth must be selected.")
+
     ### Setup ###
     emitter = _setup_emitter(click_config)
 
@@ -669,6 +679,36 @@ def preallocation(click_config,
         paint_preallocation_status(emitter=emitter,
                                    token_agent=STAKEHOLDER.token_agent,
                                    preallocation_agent=STAKEHOLDER.preallocation_escrow_agent)
+        return
+
+    # Authenticated actions: withdraw-tokens, withdraw-eth
+
+    client_account, staking_address = handle_client_account_for_staking(emitter=emitter,
+                                                                        stakeholder=STAKEHOLDER,
+                                                                        staking_address=staking_address,
+                                                                        individual_allocation=STAKEHOLDER.individual_allocation,
+                                                                        force=force)
+
+    password = None
+    if not hw_wallet and not blockchain.client.is_local:
+        password = get_client_password(checksum_address=client_account)
+
+    STAKEHOLDER.assimilate(checksum_address=client_account, password=password)
+    if withdraw_tokens:
+        token_balance = NU.from_nunits(STAKEHOLDER.token_agent.get_balance(staking_address))
+        locked_tokens = NU.from_nunits(STAKEHOLDER.preallocation_escrow_agent.unvested_tokens)
+        unlocked_tokens = token_balance - locked_tokens
+
+        emitter.echo(message=f'Collecting {unlocked_tokens} from PreallocationEscrow contract {staking_address}...')
+        receipt = STAKEHOLDER.withdraw_preallocation_tokens(unlocked_tokens)
+    elif withdraw_eth:
+        eth_balance = Web3.fromWei(blockchain.client.get_balance(staking_address), 'ether')
+        emitter.echo(message=f'Collecting {eth_balance} ETH from PreallocationEscrow contract {staking_address}...')
+        receipt = STAKEHOLDER.withdraw_preallocation_eth()
+
+    paint_receipt_summary(receipt=receipt,
+                          chain_name=STAKEHOLDER.wallet.blockchain.client.chain_name,
+                          emitter=emitter)
 
 
 def _setup_emitter(click_config):