1#!/usr/bin/env python3
2#
3# Copyright (c) 2023 Intel Corporation
4#
5# SPDX-License-Identifier: Apache-2.0
6
7'''
8This script allows flashing a mec172xevb_assy6906 board
9attached to a remote system.
10
11Usage:
12  west flash -r misc-flasher -- mec172x_remote_flasher.py <remote host>
13
14Note:
151. SSH access to remote host with write access to remote /tmp.
16   Since the script does multiple SSH connections, it is a good idea
17   to setup public key authentication and ssh-agent.
182. Dediprog "dpcmd" available in path on remote host.
19   (Can be compiled from https://github.com/DediProgSW/SF100Linux)
203. SSH user must have permission to access USB devices,
21   since dpcmd needs USB access to communicate with
22   the Dediprog programmer attached to remote host.
23
24To use with twister, a hardware map file is needed.
25Here is a sample map file:
26
27  - connected: true
28    available: true
29    id: mec172xevb_assy6906
30    platform: mec172xevb_assy6906
31    product: mec172xevb_assy6906
32    runner: misc-flasher
33    runner_params:
34      - <ZEPHYR_BASE>/boards/microchip/mec172xevb_assy6906/support/mec172x_remote_flasher.py
35      - <remote host>
36    serial_pty: "nc,<remote host>,<ser2net port>"
37
38The sample map file assumes the serial console is exposed via ser2net,
39and that it can be accessed using nc (netcat).
40
41To use twister:
42  ./scripts/twister --hardware-map <hw map file> --device-testing
43
44Required:
45* Fabric (https://www.fabfile.org/)
46'''
47
48import argparse
49import hashlib
50import pathlib
51import sys
52
53from datetime import datetime
54
55import fabric
56from invoke.exceptions import UnexpectedExit
57
58def calc_sha256(spi_file):
59    '''
60    Calculate a SHA256 of the SPI binary content plus current
61    date string.
62
63    This is used for remote file name to avoid file name
64    collision.
65    '''
66    sha256 = hashlib.sha256()
67
68    # Use SPI file content to calculate SHA.
69    with open(spi_file, "rb") as fbin:
70        spi_data = fbin.read()
71        sha256.update(spi_data)
72
73    # Add a date/time to SHA to hopefully
74    # further avoid file name collision.
75    now = datetime.now().isoformat()
76    sha256.update(now.encode("utf-8"))
77
78    return sha256.hexdigest()
79
80def parse_args():
81    '''
82    Parse command line arguments.
83    '''
84    parser = argparse.ArgumentParser(allow_abbrev=False)
85
86    # Fixed arguments
87    parser.add_argument("build_dir",
88                        help="Build directory")
89    parser.add_argument("remote_host",
90                        help="Remote host name or IP address")
91
92    # Arguments about remote machine
93    remote = parser.add_argument_group("Remote Machine")
94    remote.add_argument("--remote-tmp", required=False,
95                        help="Remote temporary directory to store SPI binary "
96                             "[default=/tmp for Linux remote]")
97    remote.add_argument("--dpcmd", required=False, default="dpcmd",
98                        help="Full path to dpcmd on remote machine")
99
100    # Remote machine type.
101    # This affects how remote path is constructed.
102    remote_type = remote.add_mutually_exclusive_group()
103    remote_type.add_argument("--remote-is-linux", required=False,
104                             default=True, action="store_true",
105                             help="Set if remote machine is a Linux-like machine [default]")
106    remote_type.add_argument("--remote-is-win", required=False,
107                             action="store_true",
108                             help="Set if remote machine is a Windows machine")
109
110    return parser.parse_args()
111
112def main():
113    '''
114    Main
115    '''
116    args = parse_args()
117
118    # Check for valid arguments and setup variables.
119    if not args.remote_tmp:
120        if args.remote_is_win:
121            # Do not assume a default temporary on Windows,
122            # as it is usually under user's directory and
123            # we do not know enough to construct a valid path
124            # at this time.
125            print("[ERROR] --remote-tmp is required for --remote-is-win")
126            sys.exit(1)
127
128        if args.remote_is_linux:
129            remote_tmp = pathlib.PurePosixPath("/tmp")
130    else:
131        if args.remote_is_win:
132            remote_tmp = pathlib.PureWindowsPath(args.remote_tmp)
133        elif args.remote_is_linux:
134            remote_tmp = pathlib.PurePosixPath(args.remote_tmp)
135
136    # Construct full path to SPI binary.
137    spi_file_path = pathlib.Path(args.build_dir)
138    spi_file_path = spi_file_path.joinpath("zephyr", "spi_image.bin")
139
140    # Calculate a sha256 digest for SPI file.
141    # This is used for remote file to avoid file name collision
142    # if there are multiple MEC17x attached to remote machine
143    # and all are trying to flash at same time.
144    sha256 = calc_sha256(spi_file_path)
145
146    # Construct full path on remote to store
147    # the transferred SPI binary.
148    remote_file_name = remote_tmp.joinpath(f"mec172x_{sha256}.bin")
149
150    print(f"[INFO] Build directory: {args.build_dir}")
151    print(f"[INFO] Remote host: {args.remote_host}")
152
153    # Connect to remote host via SSH.
154    ssh = fabric.Connection(args.remote_host, forward_agent=True)
155
156    print("[INFO] Sending file...")
157    print(f"[INFO]   Local SPI file: {spi_file_path}")
158    print(f"[INFO]   Remote SPI file: {remote_file_name}")
159
160    # Open SFTP channel, and send the SPI binary over.
161    sftp = ssh.sftp()
162    sftp.put(str(spi_file_path), str(remote_file_name))
163
164    # Run dpcmd to flash the device.
165    try:
166        dpcmd_cmd = f"{args.dpcmd} --auto {str(remote_file_name)} --verify"
167        print(f"[INFO] Invoking: {dpcmd_cmd}...")
168        ssh.run(dpcmd_cmd)
169    except UnexpectedExit:
170        print("[ERR ] Cannot flashing SPI binary!")
171
172    # Remove temporary file.
173    print(f"[INFO] Removing remote file {remote_file_name}")
174    sftp.remove(str(remote_file_name))
175
176    sftp.close()
177    ssh.close()
178
179if __name__ == "__main__":
180    main()
181