1# This Source Code Form is subject to the terms of the Mozilla Public
2# License, v. 2.0. If a copy of the MPL was not distributed with this
3# file, You can obtain one at http://mozilla.org/MPL/2.0/.
4from __future__ import absolute_import
5
6import argparse
7import os
8import signal
9import sys
10
11import mozinfo
12import mozlog.commandline
13
14from . import get_playback
15from .utils import LOG, TOOLTOOL_PATHS
16
17EXIT_SUCCESS = 0
18EXIT_EARLY_TERMINATE = 3
19EXIT_EXCEPTION = 4
20
21
22def main():
23    parser = argparse.ArgumentParser()
24    parser.add_argument(
25        "--mode",
26        default="playback",
27        choices=["record", "playback"],
28        help="Proxy server mode. Use `playback` to replay from the provided file(s). "
29        "Use `record` to generate a new recording at the path specified by `--file`. "
30        "playback - replay from provided file. "
31        "record - generate a new recording at the specified path.",
32    )
33    parser.add_argument(
34        "file",
35        nargs="+",
36        help="The playback files to replay, or the file that a recording will be saved to. "
37        "For playback, it can be any combination of the following: zip file, manifest file, "
38        "or a URL to zip/manifest file. "
39        "For recording, it's a zip fle.",
40    )
41    parser.add_argument(
42        "--host", default="localhost", help="The host to use for the proxy server."
43    )
44    parser.add_argument(
45        "--tool",
46        default="mitmproxy",
47        help="The playback tool to use (default: %(default)s).",
48    )
49    parser.add_argument(
50        "--tool-version",
51        default="6.0.2",
52        help="The playback tool version to use (default: %(default)s)",
53    )
54    parser.add_argument(
55        "--app", default="firefox", help="The app being tested (default: %(default)s)."
56    )
57    parser.add_argument(
58        "--binary",
59        required=True,
60        help=("The path to the binary being tested (typically firefox)."),
61    )
62    parser.add_argument(
63        "--topsrcdir",
64        required=True,
65        help="The top of the source directory for this project.",
66    )
67    parser.add_argument(
68        "--objdir", required=True, help="The object directory for this build."
69    )
70    parser.add_argument(
71        "--profiledir", default=None, help="Path to the profile directory."
72    )
73    parser.add_argument(
74        "--local",
75        action="store_true",
76        help="Run this locally (i.e. not in production).",
77    )
78    parser.add_argument(
79        "--verbose",
80        action="store_true",
81        default=False,
82        help="Run this locally (i.e. not in production).",
83    )
84
85    mozlog.commandline.add_logging_group(parser)
86
87    args = parser.parse_args()
88    mozlog.commandline.setup_logging("mozproxy", args, {"raw": sys.stdout})
89
90    TOOLTOOL_PATHS.append(
91        os.path.join(
92            args.topsrcdir, "python", "mozbuild", "mozbuild", "action", "tooltool.py"
93        )
94    )
95
96    if hasattr(signal, "SIGBREAK"):
97        # Terminating on windows is slightly different than other platforms.
98        # On POSIX, we just let Python's default SIGINT handler raise a
99        # KeyboardInterrupt. This doesn't work on Windows, so instead we wait
100        # for a Ctrl+Break event and raise our own KeyboardInterrupt.
101        def handle_sigbreak(sig, frame):
102            raise KeyboardInterrupt()
103
104        signal.signal(signal.SIGBREAK, handle_sigbreak)
105
106    try:
107        if args.mode == "playback":
108            if len(args.file) == 0:
109                raise Exception("Please provide at least one recording file!")
110
111            # Playback mode
112            proxy_service = get_playback(
113                {
114                    "run_local": args.local,
115                    "host": args.host,
116                    "binary": args.binary,
117                    "obj_path": args.objdir,
118                    "platform": mozinfo.os,
119                    "playback_tool": args.tool,
120                    "playback_version": args.tool_version,
121                    "playback_files": args.file,
122                    "app": args.app,
123                    "local_profile_dir": args.profiledir,
124                    "verbose": args.verbose,
125                }
126            )
127        if args.mode == "record":
128            # Record mode
129            if len(args.file) > 1:
130                raise Exception("Please provide only one recording file!")
131
132            LOG.info("Recording will be saved to: %s" % args.file)
133            proxy_service = get_playback(
134                {
135                    "run_local": args.local,
136                    "host": args.host,
137                    "binary": args.binary,
138                    "obj_path": args.objdir,
139                    "platform": mozinfo.os,
140                    "playback_tool": args.tool,
141                    "playback_version": args.tool_version,
142                    "record": True,
143                    "recording_file": args.file[0],
144                    "app": args.app,
145                    "local_profile_dir": args.profiledir,
146                    "verbose": args.verbose,
147                }
148            )
149        LOG.info("Proxy settings %s" % proxy_service)
150        proxy_service.start()
151        LOG.info("Proxy running on port %d" % proxy_service.port)
152        # Wait for a keyboard interrupt from the caller so we know when to
153        # terminate.
154        proxy_service.wait()
155        return EXIT_EARLY_TERMINATE
156    except KeyboardInterrupt:
157        LOG.info("Terminating mozproxy")
158        proxy_service.stop()
159        return EXIT_SUCCESS
160    except Exception as e:
161        LOG.error(str(e), exc_info=True)
162        return EXIT_EXCEPTION
163
164
165if __name__ == "__main__":
166    main()
167