Coverage for pyngrok/process.py: 94.21%
190 statements
« prev ^ index » next coverage.py v7.3.2, created at 2023-12-04 15:11 +0000
« prev ^ index » next coverage.py v7.3.2, created at 2023-12-04 15:11 +0000
1import atexit
2import logging
3import os
4import subprocess
5import threading
6import time
7from http import HTTPStatus
8from typing import Dict, List, Optional, Any
9from urllib.request import Request, urlopen
11import yaml
13from pyngrok import conf, installer
14from pyngrok.conf import PyngrokConfig
15from pyngrok.exception import PyngrokNgrokError, PyngrokSecurityError, PyngrokError
16from pyngrok.installer import SUPPORTED_NGROK_VERSIONS
17from pyngrok.log import NgrokLog
19__author__ = "Alex Laird"
20__copyright__ = "Copyright 2023, Alex Laird"
21__version__ = "7.0.0"
23logger = logging.getLogger(__name__)
24ngrok_logger = logging.getLogger("{}.ngrok".format(__name__))
27class NgrokProcess:
28 """
29 An object containing information about the ``ngrok`` process.
30 """
32 def __init__(self,
33 proc: subprocess.Popen,
34 pyngrok_config: PyngrokConfig) -> None:
35 #: The child process that is running ``ngrok``.
36 self.proc: subprocess.Popen = proc
37 #: The ``pyngrok`` configuration to use with ``ngrok``.
38 self.pyngrok_config: PyngrokConfig = pyngrok_config
40 #: The API URL for the ``ngrok`` web interface.
41 self.api_url: Optional[str] = None
42 #: A list of the most recent logs from ``ngrok``, limited in size to ``max_logs``.
43 self.logs: List[NgrokLog] = []
44 #: If ``ngrok`` startup fails, this will be the log of the failure.
45 self.startup_error: Optional[str] = None
47 self._tunnel_started = False
48 self._client_connected = False
49 self._monitor_thread: Optional[threading.Thread] = None
51 def __repr__(self) -> str:
52 return "<NgrokProcess: \"{}\">".format(self.api_url)
54 def __str__(self) -> str: # pragma: no cover
55 return "NgrokProcess: \"{}\"".format(self.api_url)
57 @staticmethod
58 def _line_has_error(log: NgrokLog) -> bool:
59 return log.lvl in ["ERROR", "CRITICAL"]
61 def _log_startup_line(self, line: str) -> Optional[NgrokLog]:
62 """
63 Parse the given startup log line and use it to manage the startup state
64 of the ``ngrok`` process.
66 :param line: The line to be parsed and logged.
67 :return: The parsed log.
68 """
69 log = self._log_line(line)
71 if log is None:
72 return None
73 elif self._line_has_error(log):
74 self.startup_error = log.err
75 elif log.msg:
76 # Log ngrok startup states as they come in
77 if "starting web service" in log.msg and log.addr is not None:
78 self.api_url = "http://{}".format(log.addr)
79 elif "tunnel session started" in log.msg:
80 self._tunnel_started = True
81 elif "client session established" in log.msg:
82 self._client_connected = True
84 return log
86 def _log_line(self, line: str) -> Optional[NgrokLog]:
87 """
88 Parse, log, and emit (if ``log_event_callback`` in :class:`~pyngrok.conf.PyngrokConfig` is registered) the
89 given log line.
91 :param line: The line to be processed.
92 :return: The parsed log.
93 """
94 log = NgrokLog(line)
96 if log.line == "":
97 return None
99 ngrok_logger.log(getattr(logging, log.lvl), log.line)
100 self.logs.append(log)
101 if len(self.logs) > self.pyngrok_config.max_logs:
102 self.logs.pop(0)
104 if self.pyngrok_config.log_event_callback is not None:
105 self.pyngrok_config.log_event_callback(log)
107 return log
109 def healthy(self) -> bool:
110 """
111 Check whether the ``ngrok`` process has finished starting up and is in a running, healthy state.
113 :return: ``True`` if the ``ngrok`` process is started, running, and healthy.
114 """
115 if self.api_url is None or \
116 not self._tunnel_started or \
117 not self._client_connected:
118 return False
120 if not self.api_url.lower().startswith("http"):
121 raise PyngrokSecurityError("URL must start with \"http\": {}".format(self.api_url))
123 # Ensure the process is available for requests before registering it as healthy
124 request = Request("{}/api/tunnels".format(self.api_url))
125 response = urlopen(request)
126 if response.getcode() != HTTPStatus.OK:
127 return False
129 return self.proc.poll() is None
131 def _monitor_process(self) -> None:
132 thread = threading.current_thread()
134 thread.alive = True
135 while thread.alive and self.proc.poll() is None:
136 if self.proc.stdout is None:
137 logger.debug("No stdout when monitoring the process, this may or may not be an issue")
138 continue
140 self._log_line(self.proc.stdout.readline())
142 self._monitor_thread = None
144 def start_monitor_thread(self) -> None:
145 """
146 Start a thread that will monitor the ``ngrok`` process and its logs until it completes.
148 If a monitor thread is already running, nothing will be done.
149 """
150 if self._monitor_thread is None:
151 logger.debug("Monitor thread will be started")
153 self._monitor_thread = threading.Thread(target=self._monitor_process)
154 self._monitor_thread.daemon = True
155 self._monitor_thread.start()
157 def stop_monitor_thread(self) -> None:
158 """
159 Set the monitor thread to stop monitoring the ``ngrok`` process after the next log event. This will not
160 necessarily terminate the thread immediately, as the thread may currently be idle, rather it sets a flag
161 on the thread telling it to terminate the next time it wakes up.
163 This has no impact on the ``ngrok`` process itself, only ``pyngrok``'s monitor of the process and
164 its logs.
165 """
166 if self._monitor_thread is not None:
167 logger.debug("Monitor thread will be stopped")
169 self._monitor_thread.alive = False
172def set_auth_token(pyngrok_config: PyngrokConfig,
173 token: str) -> None:
174 """
175 Set the ``ngrok`` auth token in the config file, enabling authenticated features (for instance,
176 more concurrent tunnels, custom subdomains, etc.).
178 :param pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary.
179 :param token: The auth token to set.
180 """
181 if pyngrok_config.ngrok_version == "v2":
182 start = [pyngrok_config.ngrok_path, "authtoken", token, "--log=stdout"]
183 elif pyngrok_config.ngrok_version == "v3":
184 start = [pyngrok_config.ngrok_path, "config", "add-authtoken", token, "--log=stdout"]
185 else:
186 raise PyngrokError("\"ngrok_version\" must be a supported version: {}".format(SUPPORTED_NGROK_VERSIONS))
188 if pyngrok_config.config_path:
189 logger.info("Updating authtoken for \"config_path\": {}".format(pyngrok_config.config_path))
190 start.append("--config={}".format(pyngrok_config.config_path))
191 else:
192 logger.info(
193 "Updating authtoken for default \"config_path\" of \"ngrok_path\": {}".format(pyngrok_config.ngrok_path))
195 result = str(subprocess.check_output(start))
197 if "Authtoken saved" not in result:
198 raise PyngrokNgrokError("An error occurred when saving the auth token: {}".format(result))
201def is_process_running(ngrok_path: str) -> bool:
202 """
203 Check if the ``ngrok`` process is currently running.
205 :param ngrok_path: The path to the ``ngrok`` binary.
206 :return: ``True`` if ``ngrok`` is running from the given path.
207 """
208 if ngrok_path in _current_processes:
209 # Ensure the process is still running and hasn't been killed externally, otherwise cleanup
210 if _current_processes[ngrok_path].proc.poll() is None:
211 return True
212 else:
213 logger.debug(
214 "Removing stale process for \"ngrok_path\" {}".format(ngrok_path))
216 _current_processes.pop(ngrok_path, None)
218 return False
221def get_process(pyngrok_config: PyngrokConfig) -> NgrokProcess:
222 """
223 Get the current ``ngrok`` process for the given config's ``ngrok_path``.
225 If ``ngrok`` is not running, calling this method will first start a process with
226 :class:`~pyngrok.conf.PyngrokConfig`.
228 :param pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary.
229 :return: The ``ngrok`` process.
230 """
231 if is_process_running(pyngrok_config.ngrok_path):
232 return _current_processes[pyngrok_config.ngrok_path]
234 return _start_process(pyngrok_config)
237def kill_process(ngrok_path: str) -> None:
238 """
239 Terminate the ``ngrok`` processes, if running, for the given path. This method will not block, it will just
240 issue a kill request.
242 :param ngrok_path: The path to the ``ngrok`` binary.
243 """
244 if is_process_running(ngrok_path):
245 ngrok_process = _current_processes[ngrok_path]
247 logger.info("Killing ngrok process: {}".format(ngrok_process.proc.pid))
249 try:
250 ngrok_process.proc.kill()
251 ngrok_process.proc.wait()
252 except OSError as e: # pragma: no cover
253 # If the process was already killed, nothing to do but cleanup state
254 if e.errno != 3:
255 raise e
257 _current_processes.pop(ngrok_path, None)
258 else:
259 logger.debug("\"ngrok_path\" {} is not running a process".format(ngrok_path))
262def run_process(ngrok_path: str, args: List[str]) -> None:
263 """
264 Start a blocking ``ngrok`` process with the binary at the given path and the passed args.
266 This method is meant for invoking ``ngrok`` directly (for instance, from the command line) and is not
267 necessarily compatible with non-blocking API methods. For that, use :func:`~pyngrok.process.get_process`.
269 :param ngrok_path: The path to the ``ngrok`` binary.
270 :param args: The args to pass to ``ngrok``.
271 """
272 _validate_path(ngrok_path)
274 start = [ngrok_path] + args
275 subprocess.call(start)
278def capture_run_process(ngrok_path: str, args: List[str]) -> str:
279 """
280 Start a blocking ``ngrok`` process with the binary at the given path and the passed args. When the process
281 returns, so will this method, and the captured output from the process along with it.
283 This method is meant for invoking ``ngrok`` directly (for instance, from the command line) and is not
284 necessarily compatible with non-blocking API methods. For that, use :func:`~pyngrok.process.get_process`.
286 :param ngrok_path: The path to the ``ngrok`` binary.
287 :param args: The args to pass to ``ngrok``.
288 :return: The output from the process.
289 """
290 _validate_path(ngrok_path)
292 start = [ngrok_path] + args
293 output = subprocess.check_output(start)
295 return output.decode("utf-8").strip()
298def _validate_path(ngrok_path: str) -> None:
299 """
300 Validate the given path exists, is a ``ngrok`` binary, and is ready to be started, otherwise raise a
301 relevant exception.
303 :param ngrok_path: The path to the ``ngrok`` binary.
304 """
305 if not os.path.exists(ngrok_path):
306 raise PyngrokNgrokError(
307 "ngrok binary was not found. Be sure to call \"ngrok.install_ngrok()\" first for "
308 "\"ngrok_path\": {}".format(ngrok_path))
310 if ngrok_path in _current_processes:
311 raise PyngrokNgrokError("ngrok is already running for the \"ngrok_path\": {}".format(ngrok_path))
314def _validate_config(config_path: str) -> None:
315 with open(config_path, "r") as config_file:
316 config = yaml.safe_load(config_file)
318 if config is not None:
319 installer.validate_config(config)
322def _terminate_process(process: subprocess.Popen) -> None:
323 if process is None:
324 return
326 try:
327 process.terminate()
328 except OSError: # pragma: no cover
329 logger.debug("ngrok process already terminated: {}".format(process.pid))
332def _start_process(pyngrok_config: PyngrokConfig) -> NgrokProcess:
333 """
334 Start a ``ngrok`` process with no tunnels. This will start the ``ngrok`` web interface, against
335 which HTTP requests can be made to create, interact with, and destroy tunnels.
337 :param pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary.
338 :return: The ``ngrok`` process.
339 """
340 config_path = conf.get_config_path(pyngrok_config)
342 _validate_path(pyngrok_config.ngrok_path)
343 _validate_config(config_path)
345 start = [pyngrok_config.ngrok_path, "start", "--none", "--log=stdout"]
346 if pyngrok_config.config_path:
347 logger.info("Starting ngrok with config file: {}".format(pyngrok_config.config_path))
348 start.append("--config={}".format(pyngrok_config.config_path))
349 if pyngrok_config.auth_token:
350 logger.info("Overriding default auth token")
351 start.append("--authtoken={}".format(pyngrok_config.auth_token))
352 if pyngrok_config.region:
353 logger.info("Starting ngrok in region: {}".format(pyngrok_config.region))
354 start.append("--region={}".format(pyngrok_config.region))
356 popen_kwargs: Dict[str, Any] = {"stdout": subprocess.PIPE, "universal_newlines": True}
357 if os.name == "posix":
358 popen_kwargs.update(start_new_session=pyngrok_config.start_new_session)
359 elif pyngrok_config.start_new_session:
360 logger.warning("Ignoring start_new_session=True, which requires POSIX")
361 proc = subprocess.Popen(start, **popen_kwargs)
362 atexit.register(_terminate_process, proc)
364 logger.debug("ngrok process starting with PID: {}".format(proc.pid))
366 ngrok_process = NgrokProcess(proc, pyngrok_config)
367 _current_processes[pyngrok_config.ngrok_path] = ngrok_process
369 timeout = time.time() + pyngrok_config.startup_timeout
370 while time.time() < timeout:
371 if proc.stdout is None:
372 logger.debug("No stdout when starting the process, this may or may not be an issue")
373 break
375 line = proc.stdout.readline()
376 ngrok_process._log_startup_line(line)
378 if ngrok_process.healthy():
379 logger.debug("ngrok process has started with API URL: {}".format(ngrok_process.api_url))
381 ngrok_process.startup_error = None
383 if pyngrok_config.monitor_thread:
384 ngrok_process.start_monitor_thread()
386 break
387 elif ngrok_process.proc.poll() is not None:
388 break
390 if not ngrok_process.healthy():
391 # If the process did not come up in a healthy state, clean up the state
392 kill_process(pyngrok_config.ngrok_path)
394 if ngrok_process.startup_error is not None:
395 raise PyngrokNgrokError("The ngrok process errored on start: {}.".format(ngrok_process.startup_error),
396 ngrok_process.logs,
397 ngrok_process.startup_error)
398 else:
399 raise PyngrokNgrokError("The ngrok process was unable to start.", ngrok_process.logs)
401 return ngrok_process
404_current_processes: Dict[str, NgrokProcess] = {}