1#!/usr/bin/env python3
2
3#
4# Copyright 2024, The Android Open Source Project
5#
6# Licensed under the Apache License, Version 2.0 (the "License");
7# you may not use this file except in compliance with the License.
8# You may obtain a copy of the License at
9#
10#     http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing, software
13# distributed under the License is distributed on an "AS IS" BASIS,
14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15# See the License for the specific language governing permissions and
16# limitations under the License.
17#
18
19import http.server
20import socketserver
21import json
22import re
23import urllib.parse
24from os import path
25import socket
26import argparse
27import os
28import subprocess
29import sys
30import tempfile
31import webbrowser
32import mimetypes
33import hashlib
34import shutil
35import secrets
36
37from collections import defaultdict
38
39
40def main():
41    parser = argparse.ArgumentParser(
42        "Watches a connected device for golden file updates."
43    )
44
45    parser.add_argument(
46        "--port",
47        default=find_free_port(),
48        type=int,
49        help="Port to run test at watcher web UI on.",
50    )
51    parser.add_argument(
52        "--serial",
53        default=os.environ.get("ANDROID_SERIAL"),
54        help="The ADB device serial to pull goldens from.",
55    )
56
57    parser.add_argument(
58        "--watch",
59        nargs="*",
60        action="append",
61        help="package:subdirectory where motion goldens are expected.",
62    )
63
64    parser.add_argument(
65        "--android_build_top",
66        default=os.environ.get("ANDROID_BUILD_TOP"),
67        help="The root directory of the android checkout.",
68    )
69
70    parser.add_argument(
71        "--clean",
72        default=False,
73        type=bool,
74        help="Whether to clean the golden directory on device at startup.",
75    )
76
77    parser.add_argument(
78        "--client_url",
79        default="http://motion.teams.x20web.corp.google.com/",
80        help="The URL where the client app is deployed.",
81    )
82
83    args = parser.parse_args()
84
85    if args.android_build_top is None or not os.path.exists(args.android_build_top):
86        print("ANDROID_BUILD_TOP not set. Have you sourced envsetup.sh?")
87        sys.exit(1)
88
89    serial = args.serial
90    if not serial:
91        devices_response = subprocess.run(
92            ["adb", "devices"], check=True, capture_output=True
93        ).stdout.decode("utf-8")
94        lines = [s for s in devices_response.splitlines() if s.strip()]
95
96        if len(lines) == 1:
97            print("no adb devices found")
98            sys.exit(1)
99
100        if len(lines) > 2:
101            print("multiple adb devices found, specify --serial")
102            sys.exit(1)
103
104        serial = lines[1].split("\t")[0]
105
106    adb_client = AdbClient(serial)
107    if not adb_client.run_as_root():
108        sys.exit(1)
109
110    global android_build_top
111    android_build_top = args.android_build_top
112
113    with tempfile.TemporaryDirectory() as tmpdir:
114        global golden_watcher, this_server_address
115        golden_watcher = GoldenFileWatcher(tmpdir, adb_client)
116
117        for dir in golden_watcher.list_golden_output_directories():
118            golden_watcher.add_remote_dir(dir)
119
120        if args.watch is not None:
121            for watching in args.watch:
122                parts = watching.split(":", 1)
123                package, output_dir = parts
124                if len(parts) == 2:
125                    golden_watcher.add_remote_dir(
126                        f"/data/user/0/{package}/files/{output_dir}/"
127                    )
128                else:
129                    print(f"skipping wrongly formatted watch arg [{watching}]")
130
131        if args.clean:
132            golden_watcher.clean()
133
134        this_server_address = f"http://localhost:{args.port}"
135
136        with socketserver.TCPServer(
137            ("localhost", args.port), WatchWebAppRequestHandler, golden_watcher
138        ) as httpd:
139            uiAddress = f"{args.client_url}?token={secret_token}&port={args.port}"
140            print(f"Open UI at {uiAddress}")
141            webbrowser.open(uiAddress)
142            try:
143                httpd.serve_forever()
144            except KeyboardInterrupt:
145                print("Shutting down")
146
147
148GOLDEN_ACCESS_TOKEN_HEADER = "Golden-Access-Token"
149GOLDEN_ACCESS_TOKEN_LOCATION = os.path.expanduser("~/.config/motion-golden/.token")
150
151secret_token = None
152android_build_top = None
153golden_watcher = None
154this_server_address = None
155
156
157class WatchWebAppRequestHandler(http.server.BaseHTTPRequestHandler):
158
159    def __init__(self, *args, **kwargs):
160        self.root_directory = path.abspath(path.dirname(__file__))
161        super().__init__(*args, **kwargs)
162
163    def verify_access_token(self):
164        token = self.headers.get(GOLDEN_ACCESS_TOKEN_HEADER)
165        if not token or token != secret_token:
166            self.send_response(403, "Bad authorization token!")
167            return False
168
169        return True
170
171    def do_OPTIONS(self):
172        self.send_response(200)
173        self.send_header("Allow", "GET,POST,PUT")
174        self.add_standard_headers()
175        self.end_headers()
176        self.wfile.write(b"GET,POST,PUT")
177
178    def do_GET(self):
179
180        parsed = urllib.parse.urlparse(self.path)
181
182        if parsed.path == "/service/list":
183            self.service_list_goldens()
184        elif parsed.path.startswith("/golden/"):
185            requested_file_start_index = parsed.path.find("/", len("/golden/") + 1)
186            requested_file = parsed.path[requested_file_start_index + 1 :]
187            print(requested_file)
188            self.serve_file(golden_watcher.temp_dir, requested_file)
189        else:
190            self.send_error(404)
191
192    def do_POST(self):
193        if not self.verify_access_token():
194            return
195
196        content_type = self.headers.get("Content-Type")
197
198        # refuse to receive non-json content
199        if content_type != "application/json":
200            self.send_response(400)
201            return
202
203        length = int(self.headers.get("Content-Length"))
204        message = json.loads(self.rfile.read(length))
205
206        parsed = urllib.parse.urlparse(self.path)
207        if parsed.path == "/service/refresh":
208            self.service_refresh_goldens(message["clear"])
209        else:
210            self.send_error(404)
211
212    def do_PUT(self):
213        if not self.verify_access_token():
214            return
215
216        parsed = urllib.parse.urlparse(self.path)
217        query_params = urllib.parse.parse_qs(parsed.query)
218
219        if parsed.path == "/service/update":
220            self.service_update_golden(query_params["id"][0])
221        else:
222            self.send_error(404)
223
224    def serve_file(self, root_directory, file_relative_to_root):
225        resolved_path = path.abspath(path.join(root_directory, file_relative_to_root))
226
227        print(resolved_path)
228        print(root_directory)
229
230        if path.commonprefix(
231            [resolved_path, root_directory]
232        ) == root_directory and path.isfile(resolved_path):
233            self.send_response(200)
234            self.send_header("Content-type", mimetypes.guess_type(resolved_path)[0])
235            # Accept-ranges: bytes is needed for chrome to allow seeking the
236            # video. At this time, won't handle ranges on subsequent gets,
237            # but that is likely OK given the size of these videos and that
238            # its local only.
239            self.send_header("Accept-ranges", "bytes")
240            self.add_standard_headers()
241            self.end_headers()
242            with open(resolved_path, "rb") as f:
243                self.wfile.write(f.read())
244
245        else:
246            self.send_error(404)
247
248    def service_list_goldens(self):
249        if not self.verify_access_token():
250            return
251
252        goldens_list = []
253
254        for golden in golden_watcher.cached_goldens.values():
255
256            golden_data = {}
257            golden_data["id"] = golden.id
258            golden_data["result"] = golden.result
259            golden_data["label"] = golden.golden_identifier
260            golden_data["goldenRepoPath"] = golden.golden_repo_path
261            golden_data["updated"] = golden.updated
262
263            if isinstance(golden, CachedMotionGolden):
264                golden_data["type"] = "motion"
265                golden_data["actualUrl"] = (
266                    f"{this_server_address}/golden/{golden.checksum}/{golden.local_file[len(golden_watcher.temp_dir) + 1 :]}"
267                )
268                golden_data["videoUrl"] = (
269                    f"{this_server_address}/golden/{golden.checksum}/{golden.video_location}"
270                )
271
272            elif isinstance(golden, CachedScreenshotGolden):
273                golden_data["type"] = "screenshot"
274
275                golden_data["label"] = golden.golden_identifier
276                golden_data["actualUrl"] = (
277                    f"{this_server_address}/golden/{golden.checksum}/{golden.actual_image}"
278                )
279                if golden.expected_image:
280                    golden_data["expectedUrl"] = (
281                        f"{this_server_address}/golden/{golden.checksum}/{golden.expected_image}"
282                    )
283                if golden.diff_image:
284                    golden_data["diffUrl"] = (
285                        f"{this_server_address}/golden/{golden.checksum}/{golden.diff_image}"
286                    )
287
288            else:
289                continue
290
291            goldens_list.append(golden_data)
292
293        self.send_json(goldens_list)
294
295    def service_refresh_goldens(self, clear):
296        if clear:
297            golden_watcher.clean()
298        for dir in golden_watcher.list_golden_output_directories():
299            golden_watcher.add_remote_dir(dir)
300        golden_watcher.refresh_all_golden_files()
301        self.service_list_goldens()
302
303    def service_update_golden(self, id):
304        goldens = golden_watcher.cached_goldens.values()
305        for golden in goldens:
306            if golden.id != id:
307                print("skip", golden.id)
308                continue
309
310            if isinstance(golden, CachedMotionGolden):
311                shutil.copyfile(
312                    golden.local_file,
313                    path.join(android_build_top, golden.golden_repo_path),
314                )
315
316                golden.updated = True
317                self.send_json({"result": "OK"})
318                return
319            elif isinstance(golden, CachedScreenshotGolden):
320                shutil.copyfile(
321                    path.join(golden_watcher.temp_dir, golden.actual_image),
322                    path.join(android_build_top, golden.golden_repo_path),
323                )
324
325                golden.updated = True
326                self.send_json({"result": "OK"})
327                return
328
329        self.send_error(400)
330
331    def send_json(self, data):
332        # Replace this with code that generates your JSON data
333        data_encoded = json.dumps(data).encode("utf-8")
334        self.send_response(200)
335        self.send_header("Content-type", "application/json")
336        self.add_standard_headers()
337        self.end_headers()
338        self.wfile.write(data_encoded)
339
340    def add_standard_headers(self):
341        self.send_header("Cache-Control", "no-cache, no-store, must-revalidate")
342        self.send_header("Access-Control-Allow-Origin", "*")
343        self.send_header("Access-Control-Allow-Methods", "POST, PUT, GET, OPTIONS")
344        self.send_header(
345            "Access-Control-Allow-Headers",
346            GOLDEN_ACCESS_TOKEN_HEADER + ", Content-Type, Content-Length",
347        )
348        self.send_header("Access-Control-Expose-Headers", "Winscope-Proxy-Version")
349
350
351class GoldenFileWatcher:
352
353    def __init__(self, temp_dir, adb_client):
354        self.temp_dir = temp_dir
355        self.adb_client = adb_client
356        self.remote_dirs = set()
357
358        # name -> CachedGolden
359        self.cached_goldens = {}
360
361    def add_remote_dir(self, remote_dir):
362        self.remote_dirs.add(remote_dir)
363        self.refresh_golden_files(remote_dir)
364
365    def list_golden_output_directories(self):
366        marker_name = ".motion_test_output_marker"
367        command = f"find /data/user/0/ -type f -name {marker_name}"
368
369        files = self.run_adb_command(["shell", command]).splitlines()
370        print(f"Found {len(files)} motion directories")
371
372        return [name[: -len(marker_name)] for name in files]
373
374    def clean(self):
375        self.cached_goldens = {}
376        for remote_path in self.remote_dirs:
377            self.run_adb_command(["shell", f"rm -rf {remote_path}"])
378
379    def refresh_all_golden_files(self):
380        for remote_path in self.remote_dirs:
381            self.refresh_golden_files(remote_path)
382
383    def refresh_golden_files(self, remote_dir):
384
385        updated_goldens = self.list_golden_files(remote_dir)
386
387        for golden_remote_file in updated_goldens:
388
389            local_file = self.adb_pull(golden_remote_file)
390
391            golden = None
392            if local_file.endswith(".json"):
393                golden = self.motion_golden(remote_dir, golden_remote_file, local_file)
394            elif local_file.endswith("goldResult.textproto"):
395                golden = self.screenshot_golden(
396                    remote_dir, golden_remote_file, local_file
397                )
398
399            if golden != None:
400                self.cached_goldens[golden_remote_file] = golden
401            else:
402                print(f"skipping unknonwn golden ")
403
404    def motion_golden(self, remote_dir, remote_file, local_file):
405
406        golden = CachedMotionGolden(remote_file, local_file)
407        golden.checksum = hashlib.md5(open(local_file, "rb").read()).hexdigest()
408
409        if golden.video_location:
410            self.adb_pull_image(remote_dir, golden.video_location)
411
412        return golden
413
414    def screenshot_golden(self, remote_dir, remote_file, local_file):
415        golden = CachedScreenshotGolden(remote_file, local_file)
416
417        if golden.actual_image:
418            local_actual_image = self.adb_pull_image(remote_dir, golden.actual_image)
419            golden.checksum = hashlib.md5(
420                open(local_actual_image, "rb").read()
421            ).hexdigest()
422        if golden.expected_image:
423            self.adb_pull_image(remote_dir, golden.expected_image)
424        if golden.diff_image:
425            self.adb_pull_image(remote_dir, golden.diff_image)
426
427        return golden
428
429    def list_golden_files(self, remote_dir):
430        print(f"Polling for updated goldens")
431
432        command = f"find {remote_dir} -type f \\( -name *.json -o -name *.textproto \\)"
433
434        files = self.run_adb_command(["shell", command]).splitlines()
435        print(f"Found {len(files)} files")
436
437        return files
438
439    def adb_pull(self, remote_file):
440        local_file = os.path.join(self.temp_dir, os.path.basename(remote_file))
441        self.run_adb_command(["pull", remote_file, local_file])
442        # self.run_adb_command(["shell", "rm", remote_file])
443        return local_file
444
445    def adb_pull_image(self, remote_dir, remote_file):
446        remote_path = os.path.join(remote_dir, remote_file)
447        local_path = os.path.join(self.temp_dir, remote_file)
448        os.makedirs(os.path.dirname(local_path), exist_ok=True)
449        self.run_adb_command(["pull", remote_path, local_path])
450        # self.run_adb_command(["shell", "rm", remote_path])
451        return local_path
452
453    def run_adb_command(self, args):
454        return self.adb_client.run_adb_command(args)
455
456
457class CachedGolden:
458
459    def __init__(self, remote_file, local_file):
460        self.id = hashlib.md5(remote_file.encode("utf-8")).hexdigest()
461        self.remote_file = remote_file
462        self.local_file = local_file
463        self.updated = False
464        self.checksum = "0"
465
466
467class CachedMotionGolden(CachedGolden):
468
469    def __init__(self, remote_file, local_file):
470        motion_golden_data = None
471        with open(local_file, "r") as json_file:
472            motion_golden_data = json.load(json_file)
473        metadata = motion_golden_data["//metadata"]
474
475        self.result = metadata["result"]
476        self.golden_repo_path = metadata["goldenRepoPath"]
477        self.golden_identifier = metadata["goldenIdentifier"]
478        self.video_location = None
479        if "videoLocation" in metadata:
480          self.video_location = metadata["videoLocation"]
481
482        self.test_identifier = metadata["filmstripTestIdentifier"]
483
484        with open(local_file, "w") as json_file:
485            del motion_golden_data["//metadata"]
486            json.dump(motion_golden_data, json_file, indent=2)
487
488        super().__init__(remote_file, local_file)
489
490
491class CachedScreenshotGolden(CachedGolden):
492
493    def __init__(self, remote_file, local_file):
494
495        metadata = parse_text_proto(local_file)
496
497        self.golden_repo_path = metadata["image_location_golden"]
498        self.actual_image = metadata["image_location_test"]
499
500        match = re.search(r"_actual_(.*?)\.png$", self.actual_image)
501        if match:
502            self.golden_identifier = match.group(1)
503
504        self.expected_image = metadata["image_location_reference"]
505        self.diff_image = metadata["image_location_diff"]
506        self.result = metadata["result_type"]
507
508        super().__init__(remote_file, local_file)
509
510
511class AdbClient:
512    def __init__(self, adb_serial):
513        self.adb_serial = adb_serial
514
515    def run_as_root(self):
516        root_result = self.run_adb_command(["root"])
517        if "restarting adbd as root" in root_result:
518            self.wait_for_device()
519            return True
520        if "adbd is already running as root" in root_result:
521            return True
522
523        print(f"run_as_root returned [{root_result}]")
524
525        return False
526
527    def wait_for_device(self):
528        self.run_adb_command(["wait-for-device"])
529
530    def run_adb_command(self, args):
531        command = ["adb"]
532        command += ["-s", self.adb_serial]
533        command += args
534        return subprocess.run(command, check=True, capture_output=True).stdout.decode(
535            "utf-8"
536        )
537
538
539def find_free_port():
540    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
541        s.bind(("", 0))  # Bind to a random free port provided by the OS
542        return s.getsockname()[1]  # Get the port number
543
544
545def parse_text_proto(filename):
546    data = defaultdict(dict)
547    level = 0
548
549    with open(filename, "r") as file:
550        for line in file:
551            line = line.strip()
552
553            if line.endswith("{"):
554                level += 1
555                continue
556
557            if line == "}":
558                level -= 1
559                continue
560
561            # not consuming nested messages for now
562            if not line or line.startswith("#") or level > 0:
563                continue
564
565            key, value = line.split(":", 1)
566            if not key or not value:
567                continue
568
569            key, value = key.strip(), value.strip()
570
571            if value.startswith('"') and value.endswith('"'):
572                value = value[1:-1]
573
574            data[key] = value
575
576    return data
577
578
579def get_token() -> str:
580    try:
581        with open(GOLDEN_ACCESS_TOKEN_LOCATION, "r") as token_file:
582            token = token_file.readline()
583            return token
584    except IOError:
585        token = secrets.token_hex(32)
586        os.makedirs(os.path.dirname(GOLDEN_ACCESS_TOKEN_LOCATION), exist_ok=True)
587        try:
588            with open(GOLDEN_ACCESS_TOKEN_LOCATION, "w") as token_file:
589                token_file.write(token)
590            os.chmod(GOLDEN_ACCESS_TOKEN_LOCATION, 0o600)
591        except IOError:
592            print(
593                "Unable to save persistent token {} to {}".format(
594                    token, GOLDEN_ACCESS_TOKEN_LOCATION
595                )
596            )
597        return token
598
599
600if __name__ == "__main__":
601    secret_token = get_token()
602    main()
603