From d2e0c80f7c9931cc1802b00981ddd78f64d1754d Mon Sep 17 00:00:00 2001 From: Ilyas Niyazov Date: Tue, 25 Mar 2025 11:10:30 +0300 Subject: [PATCH] [#3] Added generate proto script create container method Signed-off-by: Ilyas Niyazov --- frostfs_sdk/client/frostfs_client.py | 2 +- .../client/models/client_environment.py | 6 +- .../client/parameters/container_param.py | 2 +- .../client/parameters/create_session_param.py | 14 +++++ frostfs_sdk/client/services/container.py | 31 ++++++++-- frostfs_sdk/client/services/session.py | 61 ++++++++++--------- frostfs_sdk/client/utils/message_helper.py | 4 +- .../client/utils/request_constructor.py | 47 ++++++++++++++ frostfs_sdk/client/utils/session_cache.py | 39 ++++++++++++ frostfs_sdk/cryptography/key_extension.py | 50 ++++++++++++++- frostfs_sdk/models/dto/meta_header.py | 47 ++++++++++++++ frostfs_sdk/models/dto/owner_id.py | 1 + .../models/mappers/container_mapper.py | 24 +++++--- .../models/mappers/meta_header_mapper.py | 26 ++++++++ frostfs_sdk/models/mappers/session_mapper.py | 44 +++++++++++++ requirements.txt | 1 + 16 files changed, 348 insertions(+), 51 deletions(-) create mode 100644 frostfs_sdk/client/parameters/create_session_param.py create mode 100644 frostfs_sdk/client/utils/request_constructor.py create mode 100644 frostfs_sdk/client/utils/session_cache.py create mode 100644 frostfs_sdk/models/dto/meta_header.py create mode 100644 frostfs_sdk/models/mappers/meta_header_mapper.py create mode 100644 frostfs_sdk/models/mappers/session_mapper.py diff --git a/frostfs_sdk/client/frostfs_client.py b/frostfs_sdk/client/frostfs_client.py index d5738fa..5113587 100644 --- a/frostfs_sdk/client/frostfs_client.py +++ b/frostfs_sdk/client/frostfs_client.py @@ -1,7 +1,7 @@ # Create channel and Stubs import grpc -from frostfs_sdk.client.services.session import SessionCache +from frostfs_sdk.client.utils.session_cache import SessionCache from frostfs_sdk.client.models.client_environment import ClientEnvironment from frostfs_sdk.client.models.client_settings import ClientSettings from frostfs_sdk.client.models.ecdsa_model import ECDSA diff --git a/frostfs_sdk/client/models/client_environment.py b/frostfs_sdk/client/models/client_environment.py index 00f5aa1..29631a6 100644 --- a/frostfs_sdk/client/models/client_environment.py +++ b/frostfs_sdk/client/models/client_environment.py @@ -3,7 +3,7 @@ from frostfs_sdk.cryptography.key_extension import KeyExtension from frostfs_sdk.client.models.ecdsa_model import ECDSA from frostfs_sdk.models.dto.version import Version from frostfs_sdk.models.dto.owner_id import OwnerId -from frostfs_sdk.client.services.session import SessionCache +from frostfs_sdk.client.utils.session_cache import SessionCache class ClientEnvironment: @@ -11,10 +11,10 @@ class ClientEnvironment: self.ecdsa = ecdsa self.channel = channel self.version = version - # self.owner_id = OwnerId(KeyExtension.get_owner_id_by_wif(ecdsa.wif)) - self.owner_id = "11" + self.owner_id = OwnerId(KeyExtension().get_owner_id_by_public_key(ecdsa.public_key)) self.session_cache = session_cache self.address = address + self._session_key = None def get_session_key(self): if not self._session_key: diff --git a/frostfs_sdk/client/parameters/container_param.py b/frostfs_sdk/client/parameters/container_param.py index ed4f987..464c070 100644 --- a/frostfs_sdk/client/parameters/container_param.py +++ b/frostfs_sdk/client/parameters/container_param.py @@ -2,7 +2,7 @@ from dataclasses import dataclass, field from typing import Optional, Dict from frostfs_sdk.models.dto.container import Container -from frostfs_sdk.client.services.session import SessionToken +from frostfs_sdk.client.utils.session_cache import SessionToken from frostfs_sdk.client.parameters.wait_param import WaitParam diff --git a/frostfs_sdk/client/parameters/create_session_param.py b/frostfs_sdk/client/parameters/create_session_param.py new file mode 100644 index 0000000..a694974 --- /dev/null +++ b/frostfs_sdk/client/parameters/create_session_param.py @@ -0,0 +1,14 @@ +from dataclasses import dataclass, field +from typing import Dict, Optional + +@dataclass +class CreateSessionParam: + """ + Represents parameters for creating a session. + """ + expiration: int # -1 indicates maximum expiration + x_headers: Optional[Dict[str, str]] = field(default_factory=dict) + + def __init__(self, expiration: int, x_headers: Optional[Dict[str, str]] = None): + self.expiration = expiration + self.x_headers = x_headers if x_headers is not None else {} diff --git a/frostfs_sdk/client/services/container.py b/frostfs_sdk/client/services/container.py index 62fadbd..4d7beca 100644 --- a/frostfs_sdk/client/services/container.py +++ b/frostfs_sdk/client/services/container.py @@ -1,13 +1,18 @@ # implementation Conainer methods from frostfs_sdk.client.models.client_environment import ClientEnvironment from frostfs_sdk.client.services.context_accessor import ContextAccessor +from frostfs_sdk.client.services.session import SessionClient from frostfs_sdk.client.parameters.container_param import ContainerCreateParam from frostfs_sdk.client.parameters.call_context_param import CallContextParam +from frostfs_sdk.client.utils.session_cache import SessionToken +from frostfs_sdk.client.utils.request_constructor import RequestConstructor from frostfs_sdk.cryptography.signer import Signer from frostfs_sdk.models.dto.container import ContainerId, Container from frostfs_sdk.models.mappers.container_mapper import ContainerMapper from frostfs_sdk.models.mappers.owner_id_mapper import OwnerIdMapper from frostfs_sdk.models.mappers.version_mapper import VersionMapper +from frostfs_sdk.client.parameters.create_session_param import CreateSessionParam + from frostfs_sdk.protos.models.container import service_pb2 as service_pb2_container from frostfs_sdk.protos.models.container import service_pb2_grpc as service_pb2_grpc_container @@ -27,11 +32,8 @@ class ContainerClient(ContextAccessor): """ Creates a PUT request for creating a container. """ - grpc_container=ContainerMapper().to_grpc_message(param.container) - if not grpc_container.owner_id: - grpc_container.owner_id = OwnerIdMapper.to_grpc_message(self.get_context.owner_id) - if not grpc_container.version: - grpc_container.version = VersionMapper.to_grpc_message(self.get_context.version) + + grpc_container=ContainerMapper().to_grpc_message(param.container, self.get_context) body = service_pb2_container.PutRequest.Body( container=grpc_container, @@ -39,5 +41,24 @@ class ContainerClient(ContextAccessor): ) request = service_pb2_container.PutRequest(body=body) + session_token = self.get_or_create_session(param.session_token, ctx) + proto_token = session_token + RequestConstructor.add_meta_header(request, param.x_headers, proto_token) + signed_request = Signer.sign(self.ecdsa.private_key, request) return signed_request + + def get_or_create_session(self, session_ctx: SessionToken, ctx: CallContextParam) -> bytes: + if session_ctx: + return session_ctx.token + + session_token_from_cache = self.get_context.session_cache.try_get_value(self.get_context.get_session_key()) + if session_token_from_cache: + return session_token_from_cache.token + + new_session_token = SessionClient(self.get_context).create_session(CreateSessionParam(expiration=-1), ctx) + if new_session_token: + self.get_context.session_cache.set_value(self.get_context.get_session_key(), new_session_token) + return new_session_token.token + + raise ValueError("cannot create session") diff --git a/frostfs_sdk/client/services/session.py b/frostfs_sdk/client/services/session.py index 5b3ec5a..5d191fd 100644 --- a/frostfs_sdk/client/services/session.py +++ b/frostfs_sdk/client/services/session.py @@ -1,37 +1,38 @@ +from typing import Optional from dataclasses import dataclass +from frostfs_sdk.client.utils.session_cache import SessionToken +from frostfs_sdk.cryptography.signer import Signer +from frostfs_sdk.models.mappers.session_mapper import SessionMapper +from frostfs_sdk.models.mappers.owner_id_mapper import OwnerIdMapper +from frostfs_sdk.client.models.client_environment import ClientEnvironment +from frostfs_sdk.client.services.context_accessor import ContextAccessor +from frostfs_sdk.client.utils.request_constructor import RequestConstructor +from frostfs_sdk.client.parameters.call_context_param import CallContextParam +from frostfs_sdk.client.parameters.create_session_param import CreateSessionParam -@dataclass(frozen=True) -class SessionToken: - token: bytes +from frostfs_sdk.protos.models.session import service_pb2_grpc as service_pb2_grpc_session +from frostfs_sdk.protos.models.session import service_pb2 as service_pb2_session +from frostfs_sdk.protos.models.session import types_pb2 as types_pb2_session -class SessionCache: - def __init__(self, session_expiration_duration): - self.cache = {} - self.token_duration = session_expiration_duration - self.current_epoch = 0 +class SessionClient(ContextAccessor): + def __init__(self, client_environment: ClientEnvironment): + super().__init__(client_environment) + self.session_stub = service_pb2_grpc_session.SessionServiceStub(client_environment.channel) + def create_session(self, param: CreateSessionParam, ctx: CallContextParam) -> SessionToken: + body = service_pb2_session.CreateRequest.Body( + owner_id=OwnerIdMapper.to_grpc_message(self.get_context.owner_id), + expiration=param.expiration + ) + request = service_pb2_session.CreateRequest( + body=body + ) + RequestConstructor.add_meta_header(request, None, None) + signed_request = Signer.sign_message(self.get_context.ecdsa.private_key, request) + response: service_pb2_session.CreateResponse = self.session_stub.Create(request) - def contains(self, key): - return key in self.cache - - def try_get_value(self, key): - if not key: - return None - return self.cache.get(key) - - - def set_value(self, key, value): - if key is not None: - self.cache[key] = value - - def delete_by_prefix(self, prefix): - # Collect keys to avoid modifying dictionary during iteration - keys_to_delete = [key for key in self.cache if key.startswith(prefix)] - for key in keys_to_delete: - del self.cache[key] - - @staticmethod - def form_cache_key(address: str, key: str): - return address + key + session_token_grpc = types_pb2_session.SessionToken(response.body) + token = SessionMapper.serialize(session_token_grpc) + return SessionToken(token=token) diff --git a/frostfs_sdk/client/utils/message_helper.py b/frostfs_sdk/client/utils/message_helper.py index ca1a12a..2806109 100644 --- a/frostfs_sdk/client/utils/message_helper.py +++ b/frostfs_sdk/client/utils/message_helper.py @@ -20,11 +20,11 @@ class MessageHelper: if not field_descriptor: raise ValueError(f"Field '{field_name}' not found in message descriptor") - return message.GetField(field_descriptor[field_name]) + return getattr(field_descriptor, field_name) @staticmethod def set_field(message: Message, field_name: str, value: Any) -> None: if message is None or not field_name.strip() or value is None: raise ValueError("Some parameter is missing") - message.SetField(message.DESCRIPTOR.fields[field_name], value) + setattr(message, message.DESCRIPTOR.fields[field_name], value) diff --git a/frostfs_sdk/client/utils/request_constructor.py b/frostfs_sdk/client/utils/request_constructor.py new file mode 100644 index 0000000..ab3174b --- /dev/null +++ b/frostfs_sdk/client/utils/request_constructor.py @@ -0,0 +1,47 @@ +from google.protobuf.message import Message +from typing import Dict, Optional + +from frostfs_sdk.models.mappers.meta_header_mapper import MetaHeaderMapper +from frostfs_sdk.models.mappers.session_mapper import SessionMapper +from frostfs_sdk.models.dto.meta_header import MetaHeader +from frostfs_sdk.protos.models.session import types_pb2 as types_pb2_session + + +META_HEADER_FIELD_NAME = "meta_header" + + +class RequestConstructor: + @staticmethod + def add_meta_header(request: Message, x_headers: Optional[Dict[str, str]] = None, session_token: types_pb2_session.SessionToken = None): + """ + Adds a meta header to the request. + + :param request: A Protobuf Message.Builder object. + :param x_headers: Optional dictionary of custom headers. + :param session_token: Optional session token. + :raises ValueError: If the request or required fields are missing. + """ + if request is None: + return + + descriptor = request.DESCRIPTOR + if getattr(descriptor, META_HEADER_FIELD_NAME) is None: + raise ValueError(f"Required Protobuf field is missing: {META_HEADER_FIELD_NAME}") + + meta_header = getattr(request, META_HEADER_FIELD_NAME) + if meta_header.ByteSize() > 0: + return + + meta_header_builder = MetaHeaderMapper.to_grpc_message(MetaHeader()) + + if session_token and session_token.ByteSize() > 0: + meta_header_builder.session_token = session_token + + if x_headers: + grpc_x_headers = [ + types_pb2_session.XHeader(key=key, value=value) + for key, value in x_headers.items() + ] + meta_header_builder.x_headers.extend(grpc_x_headers) + + setattr(request, META_HEADER_FIELD_NAME, meta_header_builder.build()) diff --git a/frostfs_sdk/client/utils/session_cache.py b/frostfs_sdk/client/utils/session_cache.py new file mode 100644 index 0000000..f77b671 --- /dev/null +++ b/frostfs_sdk/client/utils/session_cache.py @@ -0,0 +1,39 @@ + +from dataclasses import dataclass +from typing import Optional + + +@dataclass(frozen=True) +class SessionToken: + token: bytes + + +class SessionCache: + def __init__(self, session_expiration_duration): + self.cache = {} + self.token_duration = session_expiration_duration + self.current_epoch = 0 + + + def contains(self, key: str): + return key in self.cache + + def try_get_value(self, key: str) -> Optional[SessionToken]: + if not key: + return None + return self.cache.get(key) + + + def set_value(self, key: str, value: SessionToken): + if key is not None: + self.cache[key] = value + + def delete_by_prefix(self, prefix: str): + # Collect keys to avoid modifying dictionary during iteration + keys_to_delete = [key for key in self.cache if key.startswith(prefix)] + for key in keys_to_delete: + del self.cache[key] + + @staticmethod + def form_cache_key(address: str, key: str): + return address + key diff --git a/frostfs_sdk/cryptography/key_extension.py b/frostfs_sdk/cryptography/key_extension.py index a6e5f13..cd8b07f 100644 --- a/frostfs_sdk/cryptography/key_extension.py +++ b/frostfs_sdk/cryptography/key_extension.py @@ -1,5 +1,18 @@ import base58 import ecdsa +import hashlib +from struct import pack, unpack +from Crypto.Hash import RIPEMD160 + + +COMPRESSED_PUBLIC_KEY_LENGTH = 33 +NEO_ADDRESS_VERSION = 0x35 +UNCOMPRESSED_PUBLIC_KEY_LENGTH = 65 +DECODE_ADDRESS_LENGTH = 21 +PS_IN_HASH160 = 0x0C +CHECK_SIG_DESCRIPTOR = int.from_bytes( + hashlib.sha256("System.Crypto.CheckSig".encode('ascii')).digest()[:4], byteorder='little' + ) class KeyExtension: @@ -37,12 +50,47 @@ class KeyExtension: compressed_public_key = public_key.to_string("compressed") return compressed_public_key + def get_owner_id_by_public_key(self, public_key: bytes) -> str: + if len(public_key) != COMPRESSED_PUBLIC_KEY_LENGTH: + raise ValueError(f"Encoded compressed public key has wrong length. Expected {COMPRESSED_PUBLIC_KEY_LENGTH}, got {len(public_key)}") + + script_hash = self.get_script_hash(public_key) + data = bytearray(DECODE_ADDRESS_LENGTH) + data[0] = NEO_ADDRESS_VERSION + data[1:] = script_hash + return base58.b58encode_check(data).decode('utf-8') + + def get_script_hash(self, public_key: bytes): + script = self.create_signature_redeem_script(public_key) + sha256_hash = hashlib.sha256(script).digest() + return self.get_ripemd160(sha256_hash) + + @staticmethod + def create_signature_redeem_script(public_key: bytes): + if len(public_key) != COMPRESSED_PUBLIC_KEY_LENGTH: + raise ValueError(f"Encoded compressed public key has wrong length. Expected {COMPRESSED_PUBLIC_KEY_LENGTH}, got {len(public_key)}") + + script = bytearray([PS_IN_HASH160, COMPRESSED_PUBLIC_KEY_LENGTH]) + script.extend(public_key) + script.append(UNCOMPRESSED_PUBLIC_KEY_LENGTH) + script.extend(pack(" Optional[types_pb2_container.Container]: + def to_grpc_message(container: Container, client_context: ClientEnvironment) -> Optional[types_pb2_container.Container]: """ Converts Container DTO to gRPC message @@ -30,18 +32,24 @@ class ContainerMapper: for k, v in container.attributes.items() ] + if container.owner_id: + owner_id = OwnerIdMapper.to_grpc_message(container.owner_id) + else: + owner_id = OwnerIdMapper.to_grpc_message(client_context.owner_id) + + if container.version: + version = VersionMapper.to_grpc_message(container.version) + else: + version = VersionMapper.to_grpc_message(client_context.version) + grpc_container = types_pb2_container.Container( nonce=container.nonce.bytes, placement_policy=PlacementPolicyMapper.to_grpc_message(container.placementPolicy), + owner_id=owner_id, + version=version, attributes=attributes ) - - if container.owner_id: - grpc_container.owner_id = OwnerIdMapper.to_grpc_message(container.owner_id) - - if container.version: - grpc_container.version = VersionMapper.to_grpc_message(container.version) - + return grpc_container @staticmethod diff --git a/frostfs_sdk/models/mappers/meta_header_mapper.py b/frostfs_sdk/models/mappers/meta_header_mapper.py new file mode 100644 index 0000000..340a19f --- /dev/null +++ b/frostfs_sdk/models/mappers/meta_header_mapper.py @@ -0,0 +1,26 @@ +from frostfs_sdk.models.mappers.version_mapper import VersionMapper +from frostfs_sdk.models.dto.meta_header import MetaHeader +from frostfs_sdk.protos.models.session import types_pb2 as types_pb2_session + + +class MetaHeaderMapper: + """ + Maps a MetaHeader object to a Protobuf RequestMetaHeader object. + """ + @staticmethod + def to_grpc_message(meta_header: MetaHeader): + """ + Converts a MetaHeader object to a Protobuf RequestMetaHeader object. + + :param meta_header: A MetaHeader object. + :return: A Protobuf RequestMetaHeader object. + :raises ValueError: If the input meta_header is None. + """ + if meta_header is None: + raise ValueError(f"Input parameter is missing: {MetaHeader.__name__}") + + return types_pb2_session.RequestMetaHeader( + version=VersionMapper.to_grpc_message(meta_header.get_version()), + epoch=meta_header.get_epoch(), + ttl=meta_header.get_ttl() + ) diff --git a/frostfs_sdk/models/mappers/session_mapper.py b/frostfs_sdk/models/mappers/session_mapper.py new file mode 100644 index 0000000..4006771 --- /dev/null +++ b/frostfs_sdk/models/mappers/session_mapper.py @@ -0,0 +1,44 @@ +from google.protobuf.message import DecodeError +from frostfs_sdk.protos.models.session import types_pb2 as types_pb2_session + + +class SessionMapper: + @staticmethod + def serialize(token: types_pb2_session.SessionToken) -> bytes: + """ + Serializes a SessionToken object into a byte array. + + :param token: A SessionToken Protobuf object. + :return: A byte array representing the serialized SessionToken. + :raises ValueError: If the input token is None. + :raises Exception: If serialization fails. + """ + if token is None: + raise ValueError(f"Input parameter is missing: {types_pb2_session.SessionToken.__name__}") + + try: + # Serialize the token to bytes + return token.SerializeToString() + except Exception as e: + raise Exception(f"Serialization failed: {str(e)}") + + @staticmethod + def deserialize_session_token(bytes_data: bytes) -> types_pb2_session.SessionToken: + """ + Deserializes a byte array into a SessionToken object. + + :param bytes_data: A byte array representing the serialized SessionToken. + :return: A SessionToken Protobuf object. + :raises ValueError: If the input byte array is None or empty. + :raises Exception: If deserialization fails. + """ + if not bytes_data: + raise ValueError(f"Input parameter is missing: {types_pb2_session.SessionToken.__name__}") + + try: + # Deserialize the byte array into a SessionToken object + session_token = types_pb2_session.SessionToken() + session_token.ParseFromString(bytes_data) + return session_token + except DecodeError as e: + raise Exception(f"Deserialization failed: {str(e)}") diff --git a/requirements.txt b/requirements.txt index 0b734ec..726ddf6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,5 +2,6 @@ base58==2.1.1 ecdsa==0.19.0 grpcio==1.70.0 grpcio-tools==1.70.0 +pycryptodome==3.22.0 pytest==8.3.4