• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python3
2# Copyright 2022 The Pigweed Authors
3#
4# Licensed under the Apache License, Version 2.0 (the "License"); you may not
5# use this file except in compliance with the License. You may obtain a copy of
6# the License at
7#
8#     https://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
12# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
13# License for the specific language governing permissions and limitations under
14# the License.
15"""Python client for pw_transfer integration test."""
16
17import logging
18import socket
19import sys
20
21from google.protobuf import text_format
22from pw_hdlc.rpc import HdlcRpcClient, default_channels
23from pw_status import Status
24import pw_transfer
25from pigweed.pw_transfer import transfer_pb2
26from pigweed.pw_transfer.integration_test import config_pb2
27
28_LOG = logging.getLogger('pw_transfer_integration_test_python_client')
29_LOG.level = logging.DEBUG
30_LOG.addHandler(logging.StreamHandler(sys.stdout))
31
32HOSTNAME: str = "localhost"
33
34
35def _main() -> int:
36    if len(sys.argv) != 2:
37        _LOG.critical("Usage: PORT")
38        return 1
39
40    # The port is passed via the command line.
41    try:
42        port = int(sys.argv[1])
43    except:
44        _LOG.critical("Invalid port specified.")
45        return 1
46
47    # Load the config from stdin.
48    try:
49        text_config = sys.stdin.buffer.read()
50        config = text_format.Parse(text_config, config_pb2.ClientConfig())
51    except Exception as e:
52        _LOG.critical("Failed to parse config file from stdin: %s", e)
53        return 1
54
55    # Open a connection to the server.
56    try:
57        rpc_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
58        rpc_socket.connect((HOSTNAME, port))
59    except:
60        _LOG.critical("Failed to connect to server at %s:%d", HOSTNAME, port)
61        return 1
62
63    # Initialize an RPC client over the socket and set up the pw_transfer manager.
64    rpc_client = HdlcRpcClient(
65        lambda: rpc_socket.recv(4096),
66        [transfer_pb2],
67        default_channels(lambda data: rpc_socket.sendall(data)),
68        lambda data: _LOG.info("%s", str(data)),
69    )
70    transfer_service = rpc_client.rpcs().pw.transfer.Transfer
71    transfer_manager = pw_transfer.Manager(
72        transfer_service,
73        default_response_timeout_s=config.chunk_timeout_ms / 1000,
74        initial_response_timeout_s=config.initial_chunk_timeout_ms / 1000,
75        max_retries=config.max_retries,
76        max_lifetime_retries=config.max_lifetime_retries,
77        default_protocol_version=pw_transfer.ProtocolVersion.LATEST,
78    )
79
80    transfer_logger = logging.getLogger('pw_transfer')
81    transfer_logger.setLevel(logging.DEBUG)
82    transfer_logger.addHandler(logging.StreamHandler(sys.stdout))
83
84    # Perform the requested transfer actions.
85    for action in config.transfer_actions:
86        protocol_version = pw_transfer.ProtocolVersion(
87            int(action.protocol_version)
88        )
89
90        # Default to the latest protocol version if none is specified.
91        if protocol_version == pw_transfer.ProtocolVersion.UNKNOWN:
92            protocol_version = pw_transfer.ProtocolVersion.LATEST
93
94        if (
95            action.transfer_type
96            == config_pb2.TransferAction.TransferType.WRITE_TO_SERVER
97        ):
98            try:
99                with open(action.file_path, 'rb') as f:
100                    data = f.read()
101            except:
102                _LOG.critical(
103                    "Failed to read input file '%s'", action.file_path
104                )
105                return 1
106
107            try:
108                transfer_manager.write(
109                    action.resource_id, data, protocol_version=protocol_version
110                )
111            except pw_transfer.client.Error as e:
112                if e.status != Status(action.expected_status):
113                    _LOG.exception(
114                        "Unexpected error encountered during write transfer"
115                    )
116                    return 1
117            except:
118                _LOG.exception("Transfer (write to server) failed")
119                return 1
120        elif (
121            action.transfer_type
122            == config_pb2.TransferAction.TransferType.READ_FROM_SERVER
123        ):
124            try:
125                data = transfer_manager.read(
126                    action.resource_id, protocol_version=protocol_version
127                )
128            except pw_transfer.client.Error as e:
129                if e.status != Status(action.expected_status):
130                    _LOG.exception(
131                        "Unexpected error encountered during read transfer"
132                    )
133                    return 1
134                continue
135            except:
136                _LOG.exception("Transfer (read from server) failed")
137                return 1
138
139            try:
140                with open(action.file_path, 'wb') as f:
141                    f.write(data)
142            except:
143                _LOG.critical(
144                    "Failed to write output file '%s'", action.file_path
145                )
146                return 1
147        else:
148            _LOG.critical("Unknown transfer type: %d", action.transfer_type)
149            return 1
150
151    _LOG.info("All transfers completed successfully")
152    return 0
153
154
155if __name__ == '__main__':
156    sys.exit(_main())
157