Coverage for pyngrok/process.py: 94.21%

1 import atexit

2 import logging

3 import os

4 import subprocess

5 import threading

6 import time

7 from http import HTTPStatus

8 from typing import Dict, List, Optional, Any

9 from urllib.request import Request, urlopen

10

11 import yaml

12

13 from pyngrok import conf, installer

14 from pyngrok.conf import PyngrokConfig

15 from pyngrok.exception import PyngrokNgrokError, PyngrokSecurityError, PyngrokError

16 from pyngrok.installer import SUPPORTED_NGROK_VERSIONS

17 from pyngrok.log import NgrokLog

18

19 __author__ = "Alex Laird"

20 __copyright__ = "Copyright 2023, Alex Laird"

21 __version__ = "7.0.0"

22

23 logger = logging.getLogger(__name__)

24 ngrok_logger = logging.getLogger("{}.ngrok".format(__name__))

25

26

27 class NgrokProcess:

28 """

29 An object containing information about the ``ngrok`` process.

30 """

31

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

39

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

46

47 self._tunnel_started = False

48 self._client_connected = False

49 self._monitor_thread: Optional[threading.Thread] = None

50

51 def __repr__(self) -> str:

52 return "<NgrokProcess: \"{}\">".format(self.api_url)

53

54 def __str__(self) -> str: # pragma: no cover

55 return "NgrokProcess: \"{}\"".format(self.api_url)

56

57 @staticmethod

58 def _line_has_error(log: NgrokLog) -> bool:

59 return log.lvl in ["ERROR", "CRITICAL"]

60

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.

65

66 :param line: The line to be parsed and logged.

67 :return: The parsed log.

68 """

69 log = self._log_line(line)

70

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

83

84 return log

85

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.

90

91 :param line: The line to be processed.

92 :return: The parsed log.

93 """

94 log = NgrokLog(line)

95

96 if log.line == "":

97 return None

98

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)

103

104 if self.pyngrok_config.log_event_callback is not None:

105 self.pyngrok_config.log_event_callback(log)

106

107 return log

108

109 def healthy(self) -> bool:

110 """

111 Check whether the ``ngrok`` process has finished starting up and is in a running, healthy state.

112

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

119

120 if not self.api_url.lower().startswith("http"):

121 raise PyngrokSecurityError("URL must start with \"http\": {}".format(self.api_url))

122

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

128

129 return self.proc.poll() is None

130

131 def _monitor_process(self) -> None:

132 thread = threading.current_thread()

133

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

139

140 self._log_line(self.proc.stdout.readline())

141

142 self._monitor_thread = None

143

144 def start_monitor_thread(self) -> None:

145 """

146 Start a thread that will monitor the ``ngrok`` process and its logs until it completes.

147

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")

152

153 self._monitor_thread = threading.Thread(target=self._monitor_process)

154 self._monitor_thread.daemon = True

155 self._monitor_thread.start()

156

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.

162

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")

168

169 self._monitor_thread.alive = False

170

171

172 def 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.).

177

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))

187

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))

194

195 result = str(subprocess.check_output(start))

196

197 if "Authtoken saved" not in result:

198 raise PyngrokNgrokError("An error occurred when saving the auth token: {}".format(result))

199

200

201 def is_process_running(ngrok_path: str) -> bool:

202 """

203 Check if the ``ngrok`` process is currently running.

204

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))

215

216 _current_processes.pop(ngrok_path, None)

217

218 return False

219

220

221 def get_process(pyngrok_config: PyngrokConfig) -> NgrokProcess:

222 """

223 Get the current ``ngrok`` process for the given config's ``ngrok_path``.

224

225 If ``ngrok`` is not running, calling this method will first start a process with

226 :class:`~pyngrok.conf.PyngrokConfig`.

227

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]

233

234 return _start_process(pyngrok_config)

235

236

237 def 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.

241

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]

246

247 logger.info("Killing ngrok process: {}".format(ngrok_process.proc.pid))

248

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

256

257 _current_processes.pop(ngrok_path, None)

258 else:

259 logger.debug("\"ngrok_path\" {} is not running a process".format(ngrok_path))

260

261

262 def 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.

265

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`.

268

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)

273

274 start = [ngrok_path] + args

275 subprocess.call(start)

276

277

278 def 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.

282

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`.

285

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)

291

292 start = [ngrok_path] + args

293 output = subprocess.check_output(start)

294

295 return output.decode("utf-8").strip()

296

297

298 def _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.

302

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))

309

310 if ngrok_path in _current_processes:

311 raise PyngrokNgrokError("ngrok is already running for the \"ngrok_path\": {}".format(ngrok_path))

312

313

314 def _validate_config(config_path: str) -> None:

315 with open(config_path, "r") as config_file:

316 config = yaml.safe_load(config_file)

317

318 if config is not None:

319 installer.validate_config(config)

320

321

322 def _terminate_process(process: subprocess.Popen) -> None:

323 if process is None:

324 return

325

326 try:

327 process.terminate()

328 except OSError: # pragma: no cover

329 logger.debug("ngrok process already terminated: {}".format(process.pid))

330

331

332 def _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.

336

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)

341

342 _validate_path(pyngrok_config.ngrok_path)

343 _validate_config(config_path)

344

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))

355

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)

363

364 logger.debug("ngrok process starting with PID: {}".format(proc.pid))

365

366 ngrok_process = NgrokProcess(proc, pyngrok_config)

367 _current_processes[pyngrok_config.ngrok_path] = ngrok_process

368

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

374

375 line = proc.stdout.readline()

376 ngrok_process._log_startup_line(line)

377

378 if ngrok_process.healthy():

379 logger.debug("ngrok process has started with API URL: {}".format(ngrok_process.api_url))

380

381 ngrok_process.startup_error = None

382

383 if pyngrok_config.monitor_thread:

384 ngrok_process.start_monitor_thread()

385

386 break

387 elif ngrok_process.proc.poll() is not None:

388 break

389

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)

393

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)

400

401 return ngrok_process

402

403

404 _current_processes: Dict[str, NgrokProcess] = {}