add withdraw all contract

web_app/api/user.py

@@ -1,3 +1,4 @@
+
 """
 This module handles user-related API endpoints.
 """
@@ -21,6 +22,7 @@
     TelegramUserDBConnector,
     UserDBConnector,
 )
+from web_app.contract_tools.blockchain_call import CLIENT
 
 logger = logging.getLogger(__name__)
 router = APIRouter()  # Initialize the router
@@ -298,3 +300,35 @@ async def allow_notification(
         raise HTTPException(status_code=404, detail=str(e))
     except Exception as e:
         raise HTTPException(status_code=500, detail="Internal server error")
+
+
+@router.post(
+    "/api/withdraw-all",
+    tags=["User Operations"],
+    summary="Withdraw all tokens from user's contract",
+    response_description="Status of withdrawal operations",
+)
+async def withdraw_all(wallet_id: str) -> dict:
+    """
+    Withdraws all supported tokens from the user's contract.
+
+    :param wallet_id: The wallet ID of the user.
+    :return: detail: "Successfully initiated withdrawals for all tokens"
+    """
+    try:
+        # Get user's contract address
+        contract_address = await get_user_contract(wallet_id)
+        if not contract_address:
+            raise HTTPException(status_code=404, detail="Contract not found")
+
+        # Perform withdrawals
+        await CLIENT.withdraw_all(contract_address)
+
+        return {"detail": "Successfully initiated withdrawals for all tokens"}
+
+    except Exception as e:
+        logger.error(f"Error in withdraw_all: {e}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to withdraw tokens: {str(e)}"
+        )

web_app/contract_tools/blockchain_call.py

@@ -352,5 +352,38 @@ async def add_extra_deposit(self, contract_address: str, token_address: str, amo
             calldata=[self._convert_address(token_address), amount],
         )
 
+    async def withdraw_all(self, contract_address: str) -> List[Any]:
+        """
+        Withdraws all supported tokens from the contract by calling withdraw with amount=0.
+
+        :param contract_address: The contract address to withdraw from
+        :return: List of responses from withdraw calls. Failed withdrawals return None.
+        """
+        contract_addr_int = self._convert_address(contract_address)
+        tasks = []
+
+        for token in TokenParams.tokens():
+            try:
+                token_addr_int = self._convert_address(token.address)
+                tasks.append(
+                    self._func_call(
+                        addr=contract_addr_int,
+                        selector="withdraw",
+                        calldata=[token_addr_int, 0]
+                    )
+                )
+            except Exception as e:
+                logger.error(f"Error preparing withdrawal for token {token.address}: {str(e)}")
+
+        results = []
+        try:
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+
+            results = [None if isinstance(r, Exception) else r for r in results]
+        except Exception as e:
+            logger.error(f"Error during batch withdrawal execution: {str(e)}")
+
+        return results
+
 
 CLIENT = StarknetClient()