1#!/usr/bin/env python3
2#
3#   Copyright 2022 - The Android Open Source Project
4#
5#   Licensed under the Apache License, Version 2.0 (the "License");
6#   you may not use this file except in compliance with the License.
7#   You may obtain a copy of the License at
8#
9#       http://www.apache.org/licenses/LICENSE-2.0
10#
11#   Unless required by applicable law or agreed to in writing, software
12#   distributed under the License is distributed on an "AS IS" BASIS,
13#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14#   See the License for the specific language governing permissions and
15#   limitations under the License.
16
17import json
18import os
19import shutil
20import socket
21import subprocess
22import tarfile
23import tempfile
24
25from dataclasses import dataclass
26from datetime import datetime
27from typing import TextIO, List, Optional
28
29from acts import context
30from acts import logger
31from acts import signals
32from acts import utils
33
34from acts.controllers.fuchsia_lib.ssh import FuchsiaSSHError, SSHProvider
35from acts.controllers.fuchsia_lib.utils_lib import wait_for_port
36from acts.tracelogger import TraceLogger
37
38DEFAULT_FUCHSIA_REPO_NAME = "fuchsia.com"
39PM_SERVE_STOP_TIMEOUT_SEC = 5
40
41
42class PackageServerError(signals.TestAbortClass):
43    pass
44
45
46def random_port() -> int:
47    s = socket.socket()
48    s.bind(('', 0))
49    return s.getsockname()[1]
50
51
52@dataclass
53class Route:
54    """Represent a route in the routing table."""
55    preferred_source: Optional[str]
56
57
58def find_routes_to(dest_ip) -> List[Route]:
59    """Find the routes used to reach a destination.
60
61    Look through the routing table for the routes that would be used without
62    sending any packets. This is especially helpful for when the device is
63    currently unreachable.
64
65    Only natively supported on Linux. MacOS has iproute2mac, but it doesn't
66    support JSON formatted output.
67
68    TODO(http://b/238924195): Add support for MacOS.
69
70    Args:
71        dest_ip: IP address of the destination
72
73    Throws:
74        CalledProcessError: if the ip command returns a non-zero exit code
75        JSONDecodeError: if the ip command doesn't return JSON
76
77    Returns:
78        Routes with destination to dest_ip.
79    """
80    resp = subprocess.run(f"ip -json route get {dest_ip}".split(),
81                          capture_output=True,
82                          check=True)
83    routes = json.loads(resp.stdout)
84    return [Route(r.get("prefsrc")) for r in routes]
85
86
87def find_host_ip(device_ip: str) -> str:
88    """Find the host's source IP used to reach a device.
89
90    Not all host interfaces can talk to a given device. This limitation can
91    either be physical through hardware or virtual through routing tables.
92    Look through the routing table without sending any packets then return the
93    preferred source IP address.
94
95    Args:
96        device_ip: IP address of the device
97
98    Raises:
99        PackageServerError: if there are multiple or no routes to device_ip, or
100            if the route doesn't contain "prefsrc"
101
102    Returns:
103        The host IP used to reach device_ip.
104    """
105    routes = find_routes_to(device_ip)
106    if len(routes) != 1:
107        raise PackageServerError(
108            f"Expected only one route to {device_ip}, got {routes}")
109
110    route = routes[0]
111    if not route.preferred_source:
112        raise PackageServerError(f'Route does not contain "prefsrc": {route}')
113    return route.preferred_source
114
115
116class PackageServer:
117    """Package manager for Fuchsia; an interface to the "pm" CLI tool."""
118
119    def __init__(self, packages_archive_path: str) -> None:
120        """
121        Args:
122            packages_archive_path: Path to an archive containing the pm binary
123                and amber-files.
124        """
125        self.log: TraceLogger = logger.create_tagged_trace_logger("pm")
126
127        self._server_log: Optional[TextIO] = None
128        self._server_proc: Optional[subprocess.Popen] = None
129        self._log_path: Optional[str] = None
130
131        self._tmp_dir = tempfile.mkdtemp(prefix="packages-")
132        tar = tarfile.open(packages_archive_path, "r:gz")
133        tar.extractall(self._tmp_dir)
134
135        self._binary_path = os.path.join(self._tmp_dir, "pm")
136        self._packages_path = os.path.join(self._tmp_dir, "amber-files")
137        self._port = random_port()
138
139        self._assert_repo_has_not_expired()
140
141    def clean_up(self) -> None:
142        if self._server_proc:
143            self.stop_server()
144        if self._tmp_dir:
145            shutil.rmtree(self._tmp_dir)
146
147    def _assert_repo_has_not_expired(self) -> None:
148        """Abort if the repository metadata has expired.
149
150        Raises:
151            TestAbortClass: when the timestamp.json file has expired
152        """
153        with open(f'{self._packages_path}/repository/timestamp.json',
154                  'r') as f:
155            data = json.load(f)
156            expiresAtRaw = data["signed"]["expires"]
157            expiresAt = datetime.strptime(expiresAtRaw, '%Y-%m-%dT%H:%M:%SZ')
158            if expiresAt <= datetime.now():
159                raise signals.TestAbortClass(
160                    f'{self._packages_path}/repository/timestamp.json has expired on {expiresAtRaw}'
161                )
162
163    def start(self) -> None:
164        """Start the package server.
165
166        Does not check for errors; view the log file for any errors.
167        """
168        if self._server_proc:
169            self.log.warn(
170                "Skipping to start the server since it has already been started"
171            )
172            return
173
174        pm_command = f'{self._binary_path} serve -c 2 -repo {self._packages_path} -l :{self._port}'
175
176        root_dir = context.get_current_context().get_full_output_path()
177        epoch = utils.get_current_epoch_time()
178        time_stamp = logger.normalize_log_line_timestamp(
179            logger.epoch_to_log_line_timestamp(epoch))
180        self._log_path = os.path.join(root_dir, f'pm_server.{time_stamp}.log')
181
182        self._server_log = open(self._log_path, 'a+')
183        self._server_proc = subprocess.Popen(pm_command.split(),
184                                             preexec_fn=os.setpgrp,
185                                             stdout=self._server_log,
186                                             stderr=subprocess.STDOUT)
187        try:
188            wait_for_port('127.0.0.1', self._port)
189        except TimeoutError as e:
190            if self._server_log:
191                self._server_log.close()
192            if self._log_path:
193                with open(self._log_path, 'r') as f:
194                    logs = f.read()
195            raise TimeoutError(
196                f"pm serve failed to expose port {self._port}. Logs:\n{logs}"
197            ) from e
198
199        self.log.info(f'Serving packages on port {self._port}')
200
201    def configure_device(self,
202                         ssh: SSHProvider,
203                         repo_name=DEFAULT_FUCHSIA_REPO_NAME) -> None:
204        """Configure the device to use this package server.
205
206        Args:
207            ssh: Device SSH transport channel
208            repo_name: Name of the repo to alias this package server
209        """
210        # Remove any existing repositories that may be stale.
211        try:
212            ssh.run(f'pkgctl repo rm fuchsia-pkg://{repo_name}')
213        except FuchsiaSSHError as e:
214            if 'NOT_FOUND' not in e.result.stderr:
215                raise e
216
217        # Configure the device with the new repository.
218        host_ip = find_host_ip(ssh.config.host_name)
219        repo_url = f"http://{host_ip}:{self._port}"
220        ssh.run(
221            f"pkgctl repo add url -f 2 -n {repo_name} {repo_url}/config.json")
222        self.log.info(
223            f'Added repo "{repo_name}" as {repo_url} on device {ssh.config.host_name}'
224        )
225
226    def stop_server(self) -> None:
227        """Stop the package server."""
228        if not self._server_proc:
229            self.log.warn(
230                "Skipping to stop the server since it hasn't been started yet")
231            return
232
233        self._server_proc.terminate()
234        try:
235            self._server_proc.wait(timeout=PM_SERVE_STOP_TIMEOUT_SEC)
236        except subprocess.TimeoutExpired:
237            self.log.warn(
238                f"Taking over {PM_SERVE_STOP_TIMEOUT_SEC}s to stop. Killing the server"
239            )
240            self._server_proc.kill()
241            self._server_proc.wait(timeout=PM_SERVE_STOP_TIMEOUT_SEC)
242        finally:
243            if self._server_log:
244                self._server_log.close()
245
246        self._server_proc = None
247        self._log_path = None
248        self._server_log = None
249