Coverage for pyngrok/process.py: 96.17%
209 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-11-28 22:19 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2022-11-28 22:19 +0000
1import atexit
2import logging
3import os
4import shlex
5import subprocess
6import threading
7import time
8from http import HTTPStatus
9from urllib.request import Request, urlopen
11import yaml
13from pyngrok import conf, installer
14from pyngrok.exception import PyngrokNgrokError, PyngrokSecurityError, PyngrokError
16__author__ = "Alex Laird"
17__copyright__ = "Copyright 2022, Alex Laird"
18__version__ = "5.2.0"
20from pyngrok.installer import SUPPORTED_NGROK_VERSIONS
22logger = logging.getLogger(__name__)
23ngrok_logger = logging.getLogger("{}.ngrok".format(__name__))
25_current_processes = {}
28class NgrokProcess:
29 """
30 An object containing information about the ``ngrok`` process.
32 :var proc: The child process that is running ``ngrok``.
33 :vartype proc: subprocess.Popen
34 :var pyngrok_config: The ``pyngrok`` configuration to use with ``ngrok``.
35 :vartype pyngrok_config: PyngrokConfig
36 :var api_url: The API URL for the ``ngrok`` web interface.
37 :vartype api_url: str
38 :var logs: A list of the most recent logs from ``ngrok``, limited in size to ``max_logs``.
39 :vartype logs: list[NgrokLog]
40 :var startup_error: If ``ngrok`` startup fails, this will be the log of the failure.
41 :vartype startup_error: str
42 """
44 def __init__(self, proc, pyngrok_config):
45 self.proc = proc
46 self.pyngrok_config = pyngrok_config
48 self.api_url = None
49 self.logs = []
50 self.startup_error = None
52 self._tunnel_started = False
53 self._client_connected = False
54 self._monitor_thread = None
56 def __repr__(self):
57 return "<NgrokProcess: \"{}\">".format(self.api_url)
59 def __str__(self): # pragma: no cover
60 return "NgrokProcess: \"{}\"".format(self.api_url)
62 @staticmethod
63 def _line_has_error(log):
64 return log.lvl in ["ERROR", "CRITICAL"]
66 def _log_startup_line(self, line):
67 """
68 Parse the given startup log line and use it to manage the startup state
69 of the ``ngrok`` process.
71 :param line: The line to be parsed and logged.
72 :type line: str
73 :return: The parsed log.
74 :rtype: NgrokLog
75 """
76 log = self._log_line(line)
78 if log is None:
79 return
80 elif self._line_has_error(log):
81 self.startup_error = log.err
82 elif log.msg:
83 # Log ngrok startup states as they come in
84 if "starting web service" in log.msg and log.addr is not None:
85 self.api_url = "http://{}".format(log.addr)
86 elif "tunnel session started" in log.msg:
87 self._tunnel_started = True
88 elif "client session established" in log.msg:
89 self._client_connected = True
91 return log
93 def _log_line(self, line):
94 """
95 Parse, log, and emit (if ``log_event_callback`` in :class:`~pyngrok.conf.PyngrokConfig` is registered) the
96 given log line.
98 :param line: The line to be processed.
99 :type line: str
100 :return: The parsed log.
101 :rtype: NgrokLog
102 """
103 log = NgrokLog(line)
105 if log.line == "":
106 return None
108 ngrok_logger.log(getattr(logging, log.lvl), log.line)
109 self.logs.append(log)
110 if len(self.logs) > self.pyngrok_config.max_logs:
111 self.logs.pop(0)
113 if self.pyngrok_config.log_event_callback is not None:
114 self.pyngrok_config.log_event_callback(log)
116 return log
118 def healthy(self):
119 """
120 Check whether the ``ngrok`` process has finished starting up and is in a running, healthy state.
122 :return: ``True`` if the ``ngrok`` process is started, running, and healthy.
123 :rtype: bool
124 """
125 if self.api_url is None or \
126 not self._tunnel_started or \
127 not self._client_connected:
128 return False
130 if not self.api_url.lower().startswith("http"):
131 raise PyngrokSecurityError("URL must start with \"http\": {}".format(self.api_url))
133 # Ensure the process is available for requests before registering it as healthy
134 request = Request("{}/api/tunnels".format(self.api_url))
135 response = urlopen(request)
136 if response.getcode() != HTTPStatus.OK:
137 return False
139 return self.proc.poll() is None
141 def _monitor_process(self):
142 thread = threading.current_thread()
144 thread.alive = True
145 while thread.alive and self.proc.poll() is None:
146 self._log_line(self.proc.stdout.readline())
148 self._monitor_thread = None
150 def start_monitor_thread(self):
151 """
152 Start a thread that will monitor the ``ngrok`` process and its logs until it completes.
154 If a monitor thread is already running, nothing will be done.
155 """
156 if self._monitor_thread is None:
157 logger.debug("Monitor thread will be started")
159 self._monitor_thread = threading.Thread(target=self._monitor_process)
160 self._monitor_thread.daemon = True
161 self._monitor_thread.start()
163 def stop_monitor_thread(self):
164 """
165 Set the monitor thread to stop monitoring the ``ngrok`` process after the next log event. This will not
166 necessarily terminate the thread immediately, as the thread may currently be idle, rather it sets a flag
167 on the thread telling it to terminate the next time it wakes up.
169 This has no impact on the ``ngrok`` process itself, only ``pyngrok``'s monitor of the process and
170 its logs.
171 """
172 if self._monitor_thread is not None:
173 logger.debug("Monitor thread will be stopped")
175 self._monitor_thread.alive = False
178class NgrokLog:
179 """
180 An object containing a parsed log from the ``ngrok`` process.
182 :var line: The raw, unparsed log line.
183 :vartype line: str
184 :var t: The log's ISO 8601 timestamp.
185 :vartype t: str
186 :var lvl: The log's level.
187 :vartype lvl: str
188 :var msg: The log's message.
189 :vartype msg: str
190 :var err: The log's error, if applicable.
191 :vartype err: str
192 :var addr: The URL, if ``obj`` is "web".
193 :vartype addr: str
194 """
196 def __init__(self, line):
197 self.line = line.strip()
198 self.t = None
199 self.lvl = "NOTSET"
200 self.msg = None
201 self.err = None
202 self.addr = None
204 for i in shlex.split(self.line):
205 if "=" not in i:
206 continue
208 key, value = i.split("=", 1)
210 if key == "lvl":
211 if not value:
212 value = self.lvl
214 value = value.upper()
215 if value == "CRIT":
216 value = "CRITICAL"
217 elif value in ["ERR", "EROR"]:
218 value = "ERROR"
219 elif value == "WARN":
220 value = "WARNING"
222 if not hasattr(logging, value):
223 value = self.lvl
225 setattr(self, key, value)
227 def __repr__(self):
228 return "<NgrokLog: t={} lvl={} msg=\"{}\">".format(self.t, self.lvl, self.msg)
230 def __str__(self): # pragma: no cover
231 attrs = [attr for attr in dir(self) if not attr.startswith("_") and getattr(self, attr) is not None]
232 attrs.remove("line")
234 return " ".join("{}=\"{}\"".format(attr, getattr(self, attr)) for attr in attrs)
237def set_auth_token(pyngrok_config, token):
238 """
239 Set the ``ngrok`` auth token in the config file, enabling authenticated features (for instance,
240 more concurrent tunnels, custom subdomains, etc.).
242 :param pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary.
243 :type pyngrok_config: PyngrokConfig
244 :param token: The auth token to set.
245 :type token: str
246 """
247 if pyngrok_config.ngrok_version == "v2":
248 start = [pyngrok_config.ngrok_path, "authtoken", token, "--log=stdout"]
249 elif pyngrok_config.ngrok_version == "v3":
250 start = [pyngrok_config.ngrok_path, "config", "add-authtoken", token, "--log=stdout"]
251 else:
252 raise PyngrokError("\"ngrok_version\" must be a supported version: {}".format(SUPPORTED_NGROK_VERSIONS))
254 if pyngrok_config.config_path:
255 logger.info("Updating authtoken for \"config_path\": {}".format(pyngrok_config.config_path))
256 start.append("--config={}".format(pyngrok_config.config_path))
257 else:
258 logger.info(
259 "Updating authtoken for default \"config_path\" of \"ngrok_path\": {}".format(pyngrok_config.ngrok_path))
261 result = subprocess.check_output(start)
263 if "Authtoken saved" not in str(result):
264 raise PyngrokNgrokError("An error occurred when saving the auth token: {}".format(result))
267def is_process_running(ngrok_path):
268 """
269 Check if the ``ngrok`` process is currently running.
271 :param ngrok_path: The path to the ``ngrok`` binary.
272 :type ngrok_path: str
273 :return: ``True`` if ``ngrok`` is running from the given path.
274 """
275 if ngrok_path in _current_processes:
276 # Ensure the process is still running and hasn't been killed externally, otherwise cleanup
277 if _current_processes[ngrok_path].proc.poll() is None:
278 return True
279 else:
280 logger.debug(
281 "Removing stale process for \"ngrok_path\" {}".format(ngrok_path))
283 _current_processes.pop(ngrok_path, None)
285 return False
288def get_process(pyngrok_config):
289 """
290 Get the current ``ngrok`` process for the given config's ``ngrok_path``.
292 If ``ngrok`` is not running, calling this method will first start a process with
293 :class:`~pyngrok.conf.PyngrokConfig`.
295 :param pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary.
296 :type pyngrok_config: PyngrokConfig
297 :return: The ``ngrok`` process.
298 :rtype: NgrokProcess
299 """
300 if is_process_running(pyngrok_config.ngrok_path):
301 return _current_processes[pyngrok_config.ngrok_path]
303 return _start_process(pyngrok_config)
306def kill_process(ngrok_path):
307 """
308 Terminate the ``ngrok`` processes, if running, for the given path. This method will not block, it will just
309 issue a kill request.
311 :param ngrok_path: The path to the ``ngrok`` binary.
312 :type ngrok_path: str
313 """
314 if is_process_running(ngrok_path):
315 ngrok_process = _current_processes[ngrok_path]
317 logger.info("Killing ngrok process: {}".format(ngrok_process.proc.pid))
319 try:
320 ngrok_process.proc.kill()
321 ngrok_process.proc.wait()
322 except OSError as e: # pragma: no cover
323 # If the process was already killed, nothing to do but cleanup state
324 if e.errno != 3:
325 raise e
327 _current_processes.pop(ngrok_path, None)
328 else:
329 logger.debug("\"ngrok_path\" {} is not running a process".format(ngrok_path))
332def run_process(ngrok_path, args):
333 """
334 Start a blocking ``ngrok`` process with the binary at the given path and the passed args.
336 This method is meant for invoking ``ngrok`` directly (for instance, from the command line) and is not
337 necessarily compatible with non-blocking API methods. For that, use :func:`~pyngrok.process.get_process`.
339 :param ngrok_path: The path to the ``ngrok`` binary.
340 :type ngrok_path: str
341 :param args: The args to pass to ``ngrok``.
342 :type args: list[str]
343 """
344 _validate_path(ngrok_path)
346 start = [ngrok_path] + args
347 subprocess.call(start)
350def capture_run_process(ngrok_path, args):
351 """
352 Start a blocking ``ngrok`` process with the binary at the given path and the passed args. When the process
353 returns, so will this method, and the captured output from the process along with it.
355 This method is meant for invoking ``ngrok`` directly (for instance, from the command line) and is not
356 necessarily compatible with non-blocking API methods. For that, use :func:`~pyngrok.process.get_process`.
358 :param ngrok_path: The path to the ``ngrok`` binary.
359 :type ngrok_path: str
360 :param args: The args to pass to ``ngrok``.
361 :type args: list[str]
362 :return: The output from the process.
363 :rtype: str
364 """
365 _validate_path(ngrok_path)
367 start = [ngrok_path] + args
368 output = subprocess.check_output(start)
370 return output.decode("utf-8").strip()
373def _validate_path(ngrok_path):
374 """
375 Validate the given path exists, is a ``ngrok`` binary, and is ready to be started, otherwise raise a
376 relevant exception.
378 :param ngrok_path: The path to the ``ngrok`` binary.
379 :type ngrok_path: str
380 """
381 if not os.path.exists(ngrok_path):
382 raise PyngrokNgrokError(
383 "ngrok binary was not found. Be sure to call \"ngrok.install_ngrok()\" first for "
384 "\"ngrok_path\": {}".format(ngrok_path))
386 if ngrok_path in _current_processes:
387 raise PyngrokNgrokError("ngrok is already running for the \"ngrok_path\": {}".format(ngrok_path))
390def _validate_config(config_path):
391 with open(config_path, "r") as config_file:
392 config = yaml.safe_load(config_file)
394 if config is not None:
395 installer.validate_config(config)
398def _terminate_process(process):
399 if process is None:
400 return
402 try:
403 process.terminate()
404 except OSError: # pragma: no cover
405 logger.debug("ngrok process already terminated: {}".format(process.pid))
408def _start_process(pyngrok_config):
409 """
410 Start a ``ngrok`` process with no tunnels. This will start the ``ngrok`` web interface, against
411 which HTTP requests can be made to create, interact with, and destroy tunnels.
413 :param pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary.
414 :type pyngrok_config: PyngrokConfig
415 :return: The ``ngrok`` process.
416 :rtype: NgrokProcess
417 """
418 if pyngrok_config.config_path is not None:
419 config_path = pyngrok_config.config_path
420 else:
421 config_path = conf.DEFAULT_NGROK_CONFIG_PATH
423 _validate_path(pyngrok_config.ngrok_path)
424 _validate_config(config_path)
426 start = [pyngrok_config.ngrok_path, "start", "--none", "--log=stdout"]
427 if pyngrok_config.config_path:
428 logger.info("Starting ngrok with config file: {}".format(pyngrok_config.config_path))
429 start.append("--config={}".format(pyngrok_config.config_path))
430 if pyngrok_config.auth_token:
431 logger.info("Overriding default auth token")
432 start.append("--authtoken={}".format(pyngrok_config.auth_token))
433 if pyngrok_config.region:
434 logger.info("Starting ngrok in region: {}".format(pyngrok_config.region))
435 start.append("--region={}".format(pyngrok_config.region))
437 popen_kwargs = {"stdout": subprocess.PIPE, "universal_newlines": True}
438 if os.name == "posix":
439 popen_kwargs.update(start_new_session=pyngrok_config.start_new_session)
440 elif pyngrok_config.start_new_session:
441 logger.warning("Ignoring start_new_session=True, which requires POSIX")
442 proc = subprocess.Popen(start, **popen_kwargs)
443 atexit.register(_terminate_process, proc)
445 logger.debug("ngrok process starting with PID: {}".format(proc.pid))
447 ngrok_process = NgrokProcess(proc, pyngrok_config)
448 _current_processes[pyngrok_config.ngrok_path] = ngrok_process
450 timeout = time.time() + pyngrok_config.startup_timeout
451 while time.time() < timeout:
452 line = proc.stdout.readline()
453 ngrok_process._log_startup_line(line)
455 if ngrok_process.healthy():
456 logger.debug("ngrok process has started with API URL: {}".format(ngrok_process.api_url))
458 ngrok_process.startup_error = None
460 if pyngrok_config.monitor_thread:
461 ngrok_process.start_monitor_thread()
463 break
464 elif ngrok_process.proc.poll() is not None:
465 break
467 if not ngrok_process.healthy():
468 # If the process did not come up in a healthy state, clean up the state
469 kill_process(pyngrok_config.ngrok_path)
471 if ngrok_process.startup_error is not None:
472 raise PyngrokNgrokError("The ngrok process errored on start: {}.".format(ngrok_process.startup_error),
473 ngrok_process.logs,
474 ngrok_process.startup_error)
475 else:
476 raise PyngrokNgrokError("The ngrok process was unable to start.", ngrok_process.logs)
478 return ngrok_process