Fix: Mypy error

src/aleph/web/controllers/storage.py

@@ -1,49 +1,36 @@
-import asyncio
 import base64
 import datetime as dt
 import functools
 import logging
 from hashlib import sha256
-from io import StringIO
-from typing import Union, Tuple, Dict, Optional
+from typing import Union, Tuple
 
 import aio_pika
 from eth_account import Account
 from eth_account.messages import encode_defunct
 
 from aleph.chains.common import get_verification_buffer
-from aleph.jobs.process_pending_messages import PendingMessageProcessor
 
 from aiohttp import web
 from aiohttp.web_request import FileField
 from aleph_message.models import ItemType
 from multidict import MultiDictProxy
-
-from aleph.chains.chain_service import ChainService, LOGGER
-from aleph.chains.nuls import NulsConnector
 from aleph.db.accessors.balances import get_total_balance
 from aleph.db.accessors.files import count_file_pins, get_file
-from aleph.db.accessors.messages import get_message_status, message_exists
-from aleph.db.connection import make_session_factory
 from aleph.db.models import PendingMessageDb
 from aleph.exceptions import AlephStorageException, UnknownHashError
-from aleph.services.p2p import init_p2p_client
-from aleph.services.storage.engine import StorageEngine
-from aleph.services.storage.fileystem_engine import FileSystemStorageEngine
-from aleph.storage import StorageService
 from aleph.toolkit import json
-from aleph.toolkit.timestamp import timestamp_to_datetime
-from aleph.types.db_session import DbSessionFactory, DbSession
+from aleph.types.db_session import DbSession
 from aleph.utils import run_in_executor, item_type_from_hash
 from aleph.web.controllers.app_state_getters import (
  get_session_factory_from_request,
- get_storage_service_from_request, get_mq_channel_from_request, get_config_from_request, get_mq_conn_from_request,
+ get_storage_service_from_request,
+ get_config_from_request,
 )
 from aleph.web.controllers.utils import multidict_proxy_to_io
 from aleph.schemas.pending_messages import BasePendingMessage
 
 logger = logging.getLogger(__name__)
-from aleph.schemas.pending_messages import parse_message
 
 MAX_FILE_SIZE = 100 * 1024 * 1024
 
@@ -111,34 +98,47 @@ async def verify_signature(message: BasePendingMessage) -> bool:
  return verified
 
 
-async def get_message_content(post_data: MultiDictProxy[Union[str, bytes, FileField]]) -> Tuple[dict, int]:
+async def get_message_content(
+ post_data: MultiDictProxy[Union[str, bytes, FileField]]
+) -> Tuple[dict, int]:
  message_bytearray = post_data.get("message", b"")
  value = post_data.get("size") or 0
- if not message_bytearray:
- return {}, int(value) # Empty dictionary if no message content
 
- message_string = message_bytearray.decode("utf-8")
- message_dict = json.loads(message_string)
- message_dict["time"] = float(message_dict["time"])
+ if isinstance(message_bytearray, bytearray):
+ message_string = message_bytearray.decode("utf-8")
+ message_dict = json.loads(message_string)
+ message_dict["time"] = float(message_dict["time"])
+ else:
+ message_dict = {}
 
- return message_dict, int(value)
+ return message_dict, int(str(value))
 
 
 async def init_mq_con(config):
  return await aio_pika.connect_robust(
- host=config.p2p.mq_host.value, port=config.rabbitmq.port.value, login=config.rabbitmq.username.value,
- password=config.rabbitmq.password.value
+ host=config.p2p.mq_host.value,
+ port=config.rabbitmq.port.value,
+ login=config.rabbitmq.username.value,
+ password=config.rabbitmq.password.value,
  )
 
 
-async def verify_and_handle_request(pending_message_db, file_io, message, size):
+async def verify_and_handle_request(
+ pending_message_db, file_io, message, size, session_factory
+):
  content = file_io.read(size)
  item_content = json.loads(message["item_content"])
  actual_item_hash = sha256(content).hexdigest()
  c_item_hash = item_content["item_hash"]
 
  is_signature = await verify_signature(message=pending_message_db)
-
+ with session_factory() as session:
+ current_balance = get_total_balance(
+ session=session, address=pending_message_db.sender
+ )
+ if current_balance < len(content):
+ output = {"status": "Payment Required"}
+ return web.json_response(output, status=402)
  if not is_signature:
  output = {"status": "Forbidden"}
  return web.json_response(output, status=403)
@@ -161,9 +161,12 @@ async def storage_add_file_with_message(request: web.Request):
  post = await request.post()
  file_io = multidict_proxy_to_io(post)
  message, size = await get_message_content(post)
- pending_message_db = PendingMessageDb.from_message_dict(message_dict=message, reception_time=dt.datetime.now(),
- fetched=True)
- is_valid_message = await verify_and_handle_request(pending_message_db, file_io, message, size)
+ pending_message_db = PendingMessageDb.from_message_dict(
+ message_dict=message, reception_time=dt.datetime.now(), fetched=True
+ )
+ is_valid_message = await verify_and_handle_request(
+ pending_message_db, file_io, message, size, session_factory
+ )
  if is_valid_message is not None:
  return is_valid_message