Coverage for pyngrok/process.py: 96.17%

1 import atexit

2 import logging

3 import os

4 import shlex

5 import subprocess

6 import threading

7 import time

8 from http import HTTPStatus

9 from urllib.request import Request, urlopen

10

11 import yaml

12

13 from pyngrok import conf, installer

14 from pyngrok.exception import PyngrokNgrokError, PyngrokSecurityError, PyngrokError

15

16 __author__ = "Alex Laird"

17 __copyright__ = "Copyright 2023, Alex Laird"

18 __version__ = "5.2.2"

19

20 from pyngrok.installer import SUPPORTED_NGROK_VERSIONS

21

22 logger = logging.getLogger(__name__)

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

24

25 _current_processes = {}

26

27

28 class NgrokProcess:

29 """

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

31

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

43

44 def __init__(self, proc, pyngrok_config):

45 self.proc = proc

46 self.pyngrok_config = pyngrok_config

47

48 self.api_url = None

49 self.logs = []

50 self.startup_error = None

51

52 self._tunnel_started = False

53 self._client_connected = False

54 self._monitor_thread = None

55

56 def __repr__(self):

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

58

59 def __str__(self): # pragma: no cover

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

61

62 @staticmethod

63 def _line_has_error(log):

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

65

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.

70

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)

77

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

90

91 return log

92

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.

97

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)

104

105 if log.line == "":

106 return None

107

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)

112

113 if self.pyngrok_config.log_event_callback is not None:

114 self.pyngrok_config.log_event_callback(log)

115

116 return log

117

118 def healthy(self):

119 """

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

121

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

129

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

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

132

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

138

139 return self.proc.poll() is None

140

141 def _monitor_process(self):

142 thread = threading.current_thread()

143

144 thread.alive = True

145 while thread.alive and self.proc.poll() is None:

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

147

148 self._monitor_thread = None

149

150 def start_monitor_thread(self):

151 """

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

153

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

158

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

160 self._monitor_thread.daemon = True

161 self._monitor_thread.start()

162

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.

168

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

174

175 self._monitor_thread.alive = False

176

177

178 class NgrokLog:

179 """

180 An object containing a parsed log from the ``ngrok`` process.

181

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

195

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

203

204 for i in shlex.split(self.line):

205 if "=" not in i:

206 continue

207

208 key, value = i.split("=", 1)

209

210 if key == "lvl":

211 if not value:

212 value = self.lvl

213

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"

221

222 if not hasattr(logging, value):

223 value = self.lvl

224

225 setattr(self, key, value)

226

227 def __repr__(self):

228 return "<NgrokLog: t={} lvl={} msg=\"{}\">".format(self.t, self.lvl, self.msg)

229

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

233

234 return " ".join("{}=\"{}\"".format(attr, getattr(self, attr)) for attr in attrs)

235

236

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

241

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

253

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

260

261 result = subprocess.check_output(start)

262

263 if "Authtoken saved" not in str(result):

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

265

266

267 def is_process_running(ngrok_path):

268 """

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

270

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

282

283 _current_processes.pop(ngrok_path, None)

284

285 return False

286

287

288 def get_process(pyngrok_config):

289 """

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

291

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

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

294

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]

302

303 return _start_process(pyngrok_config)

304

305

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

310

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]

316

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

318

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

326

327 _current_processes.pop(ngrok_path, None)

328 else:

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

330

331

332 def run_process(ngrok_path, args):

333 """

334 Start a blocking ``ngrok`` process with the binary at the given path and the passed args.

335

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

338

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)

345

346 start = [ngrok_path] + args

347 subprocess.call(start)

348

349

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

354

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

357

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)

366

367 start = [ngrok_path] + args

368 output = subprocess.check_output(start)

369

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

371

372

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

377

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

385

386 if ngrok_path in _current_processes:

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

388

389

390 def _validate_config(config_path):

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

392 config = yaml.safe_load(config_file)

393

394 if config is not None:

395 installer.validate_config(config)

396

397

398 def _terminate_process(process):

399 if process is None:

400 return

401

402 try:

403 process.terminate()

404 except OSError: # pragma: no cover

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

406

407

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

412

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

422

423 _validate_path(pyngrok_config.ngrok_path)

424 _validate_config(config_path)

425

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

436

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)

444

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

446

447 ngrok_process = NgrokProcess(proc, pyngrok_config)

448 _current_processes[pyngrok_config.ngrok_path] = ngrok_process

449

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)

454

455 if ngrok_process.healthy():

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

457

458 ngrok_process.startup_error = None

459

460 if pyngrok_config.monitor_thread:

461 ngrok_process.start_monitor_thread()

462

463 break

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

465 break

466

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)

470

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)

477

478 return ngrok_process