157 lines
5.3 KiB
Python
157 lines
5.3 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright 2022 The Pigweed Authors
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
|
|
# use this file except in compliance with the License. You may obtain a copy of
|
|
# the License at
|
|
#
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
|
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
|
# License for the specific language governing permissions and limitations under
|
|
# the License.
|
|
"""Python client for pw_transfer integration test."""
|
|
|
|
import logging
|
|
import socket
|
|
import sys
|
|
|
|
from google.protobuf import text_format
|
|
from pw_hdlc.rpc import HdlcRpcClient, default_channels
|
|
from pw_status import Status
|
|
import pw_transfer
|
|
from pigweed.pw_transfer import transfer_pb2
|
|
from pigweed.pw_transfer.integration_test import config_pb2
|
|
|
|
_LOG = logging.getLogger('pw_transfer_integration_test_python_client')
|
|
_LOG.level = logging.DEBUG
|
|
_LOG.addHandler(logging.StreamHandler(sys.stdout))
|
|
|
|
HOSTNAME: str = "localhost"
|
|
|
|
|
|
def _main() -> int:
|
|
if len(sys.argv) != 2:
|
|
_LOG.critical("Usage: PORT")
|
|
return 1
|
|
|
|
# The port is passed via the command line.
|
|
try:
|
|
port = int(sys.argv[1])
|
|
except:
|
|
_LOG.critical("Invalid port specified.")
|
|
return 1
|
|
|
|
# Load the config from stdin.
|
|
try:
|
|
text_config = sys.stdin.buffer.read()
|
|
config = text_format.Parse(text_config, config_pb2.ClientConfig())
|
|
except Exception as e:
|
|
_LOG.critical("Failed to parse config file from stdin: %s", e)
|
|
return 1
|
|
|
|
# Open a connection to the server.
|
|
try:
|
|
rpc_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
rpc_socket.connect((HOSTNAME, port))
|
|
except:
|
|
_LOG.critical("Failed to connect to server at %s:%d", HOSTNAME, port)
|
|
return 1
|
|
|
|
# Initialize an RPC client over the socket and set up the pw_transfer manager.
|
|
rpc_client = HdlcRpcClient(
|
|
lambda: rpc_socket.recv(4096),
|
|
[transfer_pb2],
|
|
default_channels(lambda data: rpc_socket.sendall(data)),
|
|
lambda data: _LOG.info("%s", str(data)),
|
|
)
|
|
transfer_service = rpc_client.rpcs().pw.transfer.Transfer
|
|
transfer_manager = pw_transfer.Manager(
|
|
transfer_service,
|
|
default_response_timeout_s=config.chunk_timeout_ms / 1000,
|
|
initial_response_timeout_s=config.initial_chunk_timeout_ms / 1000,
|
|
max_retries=config.max_retries,
|
|
max_lifetime_retries=config.max_lifetime_retries,
|
|
default_protocol_version=pw_transfer.ProtocolVersion.LATEST,
|
|
)
|
|
|
|
transfer_logger = logging.getLogger('pw_transfer')
|
|
transfer_logger.setLevel(logging.DEBUG)
|
|
transfer_logger.addHandler(logging.StreamHandler(sys.stdout))
|
|
|
|
# Perform the requested transfer actions.
|
|
for action in config.transfer_actions:
|
|
protocol_version = pw_transfer.ProtocolVersion(
|
|
int(action.protocol_version)
|
|
)
|
|
|
|
# Default to the latest protocol version if none is specified.
|
|
if protocol_version == pw_transfer.ProtocolVersion.UNKNOWN:
|
|
protocol_version = pw_transfer.ProtocolVersion.LATEST
|
|
|
|
if (
|
|
action.transfer_type
|
|
== config_pb2.TransferAction.TransferType.WRITE_TO_SERVER
|
|
):
|
|
try:
|
|
with open(action.file_path, 'rb') as f:
|
|
data = f.read()
|
|
except:
|
|
_LOG.critical(
|
|
"Failed to read input file '%s'", action.file_path
|
|
)
|
|
return 1
|
|
|
|
try:
|
|
transfer_manager.write(
|
|
action.resource_id, data, protocol_version=protocol_version
|
|
)
|
|
except pw_transfer.client.Error as e:
|
|
if e.status != Status(action.expected_status):
|
|
_LOG.exception(
|
|
"Unexpected error encountered during write transfer"
|
|
)
|
|
return 1
|
|
except:
|
|
_LOG.exception("Transfer (write to server) failed")
|
|
return 1
|
|
elif (
|
|
action.transfer_type
|
|
== config_pb2.TransferAction.TransferType.READ_FROM_SERVER
|
|
):
|
|
try:
|
|
data = transfer_manager.read(
|
|
action.resource_id, protocol_version=protocol_version
|
|
)
|
|
except pw_transfer.client.Error as e:
|
|
if e.status != Status(action.expected_status):
|
|
_LOG.exception(
|
|
"Unexpected error encountered during read transfer"
|
|
)
|
|
return 1
|
|
continue
|
|
except:
|
|
_LOG.exception("Transfer (read from server) failed")
|
|
return 1
|
|
|
|
try:
|
|
with open(action.file_path, 'wb') as f:
|
|
f.write(data)
|
|
except:
|
|
_LOG.critical(
|
|
"Failed to write output file '%s'", action.file_path
|
|
)
|
|
return 1
|
|
else:
|
|
_LOG.critical("Unknown transfer type: %d", action.transfer_type)
|
|
return 1
|
|
|
|
_LOG.info("All transfers completed successfully")
|
|
return 0
|
|
|
|
|
|
if __name__ == '__main__':
|
|
sys.exit(_main())
|