Coverage for pyngrok/process.py: 92.27%

# Coverage for pyngrok/process.py : 92.27%

## 207 statements

Hot-keys on this page

r m x p toggle line displays

j k next/prev highlighted chunk

0 (zero) top of page

1 (one) first highlighted chunk

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

15

16 __author__ = "Alex Laird"

17 __copyright__ = "Copyright 2020, Alex Laird"

18 __version__ = "5.0.0"

19

20 logger = logging.getLogger(__name__)

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

22

23 _current_processes = {}

24

25

26 class NgrokProcess:

27 """

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

29

30 :var proc: The child process that is running ``ngrok``.

31 :vartype proc: subprocess.Popen

32 :var pyngrok_config: The ``pyngrok`` configuration to use with ``ngrok``.

33 :vartype pyngrok_config: PyngrokConfig

34 :var api_url: The API URL for the ``ngrok`` web interface.

35 :vartype api_url: str

36 :var logs: A list of the most recent logs from ``ngrok``, limited in size to ``max_logs``.

37 :vartype logs: list[NgrokLog]

38 :var startup_error: If ``ngrok`` startup fails, this will be the log of the failure.

39 :vartype startup_error: str

40 """

41

42 def __init__(self, proc, pyngrok_config):

43 self.proc = proc

44 self.pyngrok_config = pyngrok_config

45

46 self.api_url = None

47 self.logs = []

48 self.startup_error = None

49

50 self._tunnel_started = False

51 self._client_connected = False

52 self._monitor_thread = None

53

54 def __repr__(self):

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

56

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

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

59

60 @staticmethod

61 def _line_has_error(log):

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

63

64 def _log_startup_line(self, line):

65 """

66 Parse the given startup log line and use it to manage the startup state

67 of the ``ngrok`` process.

68

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

70 :type line: str

71 :return: The parsed log.

72 :rtype: NgrokLog

73 """

74 log = self._log_line(line)

75

76 if log is None:

77 return

78 elif self._line_has_error(log):

79 self.startup_error = log.err

80 else:

81 # Log `ngrok` startup states as they come in

82 if "starting web service" in log.msg and log.addr is not None:

83 self.api_url = "http://{}".format(log.addr)

84 elif "tunnel session started" in log.msg:

85 self._tunnel_started = True

86 elif "client session established" in log.msg:

87 self._client_connected = True

88

89 return log

90

91 def _log_line(self, line):

92 """

93 Parse, log, and emit (if ``log_event_callback`` in :class:`~pyngrok.conf.PyngrokConfig` is registered) the

94 given log line.

95

96 :param line: The line to be processed.

97 :type line: str

98 :return: The parsed log.

99 :rtype: NgrokLog

100 """

101 log = NgrokLog(line)

102

103 if log.line == "":

104 return None

105

106 ngrok_logger.log(getattr(logging, log.lvl), log.line)

107 self.logs.append(log)

108 if len(self.logs) > self.pyngrok_config.max_logs:

109 self.logs.pop(0)

110

111 if self.pyngrok_config.log_event_callback is not None:

112 self.pyngrok_config.log_event_callback(log)

113

114 return log

115

116 def healthy(self):

117 """

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

119

120 :return: ``True`` if the ``ngrok`` process is started, running, and healthy.

121 :rtype: bool

122 """

123 if self.api_url is None or \

124 not self._tunnel_started or \

125 not self._client_connected:

126 return False

127

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

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

130

131 # Ensure the process is available for requests before registering it as healthy

132 request = Request("{}/api/tunnels".format(self.api_url))

133 response = urlopen(request)

134 if response.getcode() != HTTPStatus.OK:

135 return False

136

137 return self.proc.poll() is None and self.startup_error is None

138

139 def _monitor_process(self):

140 thread = threading.current_thread()

141

142 thread.alive = True

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

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

145

146 self._monitor_thread = None

147

148 def start_monitor_thread(self):

149 """

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

151

152 If a monitor thread is already running, nothing will be done.

153 """

154 if self._monitor_thread is None:

155 logger.debug("Monitor thread will be started")

156

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

158 self._monitor_thread.daemon = True

159 self._monitor_thread.start()

160

161 def stop_monitor_thread(self):

162 """

163 Set the monitor thread to stop monitoring the ``ngrok`` process after the next log event. This will not

164 necessarily terminate the thread immediately, as the thread may currently be idle, rather it sets a flag

165 on the thread telling it to terminate the next time it wakes up.

166

167 This has no impact on the ``ngrok`` process itself, only ``pyngrok``'s monitor of the process and

168 its logs.

169 """

170 if self._monitor_thread is not None:

171 logger.debug("Monitor thread will be stopped")

172

173 self._monitor_thread.alive = False

174

175

176 class NgrokLog:

177 """

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

179

180 :var line: The raw, unparsed log line.

181 :vartype line: str

182 :var t: The log's ISO 8601 timestamp.

183 :vartype t: str

184 :var lvl: The log's level.

185 :vartype lvl: str

186 :var msg: The log's message.

187 :vartype msg: str

188 :var err: The log's error, if applicable.

189 :vartype err: str

190 :var addr: The URL, if ``obj`` is "web".

191 :vartype addr: str

192 """

193

194 def __init__(self, line):

195 self.line = line.strip()

196 self.t = None

197 self.lvl = "NOTSET"

198 self.msg = None

199 self.err = None

200 self.addr = None

201

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

203 if "=" not in i:

204 continue

205

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

207

208 if key == "lvl":

209 if not value:

210 value = self.lvl

211

212 value = value.upper()

213 if value == "CRIT":

214 value = "CRITICAL"

215 elif value in ["ERR", "EROR"]:

216 value = "ERROR"

217 elif value == "WARN":

218 value = "WARNING"

219

220 if not hasattr(logging, value):

221 value = self.lvl

222

223 setattr(self, key, value)

224

225 def __repr__(self):

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

227

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

229 attrs = [attr for attr in dir(self) if not attr.startswith("_") and getattr(self, attr) is not None]

230 attrs.remove("line")

231

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

233

234

235 def set_auth_token(pyngrok_config, token):

236 """

237 Set the ``ngrok`` auth token in the config file, enabling authenticated features (for instance,

238 more concurrent tunnels, custom subdomains, etc.).

239

240 :param pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary.

241 :type pyngrok_config: PyngrokConfig

242 :param token: The auth token to set.

243 :type token: str

244 """

245 start = [pyngrok_config.ngrok_path, "authtoken", token, "--log=stdout"]

246 if pyngrok_config.config_path:

247 logger.info("Updating authtoken for \"config_path\": {}".format(pyngrok_config.config_path))

248 start.append("--config={}".format(pyngrok_config.config_path))

249 else:

250 logger.info(

251 "Updating authtoken for default \"config_path\" of \"ngrok_path\": {}".format(pyngrok_config.ngrok_path))

252

253 result = subprocess.check_output(start)

254

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

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

257

258

259 def is_process_running(ngrok_path):

260 """

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

262

263 :param ngrok_path: The path to the ``ngrok`` binary.

264 :type ngrok_path: str

265 :return: ``True`` if ``ngrok`` is running from the given path.

266 """

267 if ngrok_path in _current_processes:

268 # Ensure the process is still running and hasn't been killed externally, otherwise cleanup

269 if _current_processes[ngrok_path].proc.poll() is None:

270 return True

271 else:

272 logger.debug(

273 "Removing stale process for \"ngrok_path\" {}".format(ngrok_path))

274

275 _current_processes.pop(ngrok_path, None)

276

277 return False

278

279

280 def get_process(pyngrok_config):

281 """

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

283

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

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

286

287 :param pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary.

288 :type pyngrok_config: PyngrokConfig

289 :return: The ``ngrok`` process.

290 :rtype: NgrokProcess

291 """

292 if is_process_running(pyngrok_config.ngrok_path):

293 return _current_processes[pyngrok_config.ngrok_path]

294

295 return _start_process(pyngrok_config)

296

297

298 def kill_process(ngrok_path):

299 """

300 Terminate the ``ngrok`` processes, if running, for the given path. This method will not block, it will just

301 issue a kill request.

302

303 :param ngrok_path: The path to the ``ngrok`` binary.

304 :type ngrok_path: str

305 """

306 if is_process_running(ngrok_path):

307 ngrok_process = _current_processes[ngrok_path]

308

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

310

311 try:

312 ngrok_process.proc.kill()

313 ngrok_process.proc.wait()

314 except OSError as e:

315 # If the process was already killed, nothing to do but cleanup state

316 if e.errno != 3:

317 raise e

318

319 _current_processes.pop(ngrok_path, None)

320 else:

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

322

323

324 def run_process(ngrok_path, args):

325 """

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

327

328 This method is meant for invoking ``ngrok`` directly (for instance, from the command line) and is not

329 necessarily compatible with non-blocking API methods. For that, use :func:`~pyngrok.process.get_process`.

330

331 :param ngrok_path: The path to the ``ngrok`` binary.

332 :type ngrok_path: str

333 :param args: The args to pass to ``ngrok``.

334 :type args: list[str]

335 """

336 _validate_path(ngrok_path)

337

338 start = [ngrok_path] + args

339 subprocess.call(start)

340

341

342 def capture_run_process(ngrok_path, args):

343 """

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

345 returns, so will this method, and the captured output from the process along with it.

346

347 This method is meant for invoking ``ngrok`` directly (for instance, from the command line) and is not

348 necessarily compatible with non-blocking API methods. For that, use :func:`~pyngrok.process.get_process`.

349

350 :param ngrok_path: The path to the ``ngrok`` binary.

351 :type ngrok_path: str

352 :param args: The args to pass to ``ngrok``.

353 :type args: list[str]

354 :return: The output from the process.

355 :rtype: str

356 """

357 _validate_path(ngrok_path)

358

359 start = [ngrok_path] + args

360 output = subprocess.check_output(start)

361

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

363

364

365 def _validate_path(ngrok_path):

366 """

367 Validate the given path exists, is a ``ngrok`` binary, and is ready to be started, otherwise raise a

368 relevant exception.

369

370 :param ngrok_path: The path to the ``ngrok`` binary.

371 :type ngrok_path: str

372 """

373 if not os.path.exists(ngrok_path):

374 raise PyngrokNgrokError(

375 "ngrok binary was not found. Be sure to call \"ngrok.install_ngrok()\" first for "

376 "\"ngrok_path\": {}".format(ngrok_path))

377

378 if ngrok_path in _current_processes:

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

380

381

382 def _validate_config(config_path):

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

384 config = yaml.safe_load(config_file)

385

386 if config is not None:

387 installer.validate_config(config)

388

389

390 def _terminate_process(process):

391 if process is None:

392 return

393

394 try:

395 process.terminate()

396 except OSError:

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

398

399

400 def _start_process(pyngrok_config):

401 """

402 Start a ``ngrok`` process with no tunnels. This will start the ``ngrok`` web interface, against

403 which HTTP requests can be made to create, interact with, and destroy tunnels.

404

405 :param pyngrok_config: The ``pyngrok`` configuration to use when interacting with the ``ngrok`` binary.

406 :type pyngrok_config: PyngrokConfig

407 :return: The ``ngrok`` process.

408 :rtype: NgrokProcess

409 """

410 if pyngrok_config.config_path is not None:

411 config_path = pyngrok_config.config_path

412 else:

413 config_path = conf.DEFAULT_NGROK_CONFIG_PATH

414

415 _validate_path(pyngrok_config.ngrok_path)

416 _validate_config(config_path)

417

418 start = [pyngrok_config.ngrok_path, "start", "--none", "--log=stdout"]

419 if pyngrok_config.config_path:

420 logger.info("Starting ngrok with config file: {}".format(pyngrok_config.config_path))

421 start.append("--config={}".format(pyngrok_config.config_path))

422 if pyngrok_config.auth_token:

423 logger.info("Overriding default auth token")

424 start.append("--authtoken={}".format(pyngrok_config.auth_token))

425 if pyngrok_config.region:

426 logger.info("Starting ngrok in region: {}".format(pyngrok_config.region))

427 start.append("--region={}".format(pyngrok_config.region))

428

429 popen_kwargs = {"stdout": subprocess.PIPE, "universal_newlines": True}

430 if os.name == "posix":

431 popen_kwargs.update(start_new_session=pyngrok_config.start_new_session)

432 elif pyngrok_config.start_new_session:

433 logger.warning("Ignoring start_new_session=True, which requires POSIX")

434 proc = subprocess.Popen(start, **popen_kwargs)

435 atexit.register(_terminate_process, proc)

436

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

438

439 ngrok_process = NgrokProcess(proc, pyngrok_config)

440 _current_processes[pyngrok_config.ngrok_path] = ngrok_process

441

442 timeout = time.time() + pyngrok_config.startup_timeout

443 while time.time() < timeout:

444 line = proc.stdout.readline()

445 ngrok_process._log_startup_line(line)

446

447 if ngrok_process.healthy():

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

449

450 if pyngrok_config.monitor_thread:

451 ngrok_process.start_monitor_thread()

452

453 break

454 elif ngrok_process.startup_error is not None or \

455 ngrok_process.proc.poll() is not None:

456 break

457

458 if not ngrok_process.healthy():

459 # If the process did not come up in a healthy state, clean up the state

460 kill_process(pyngrok_config.ngrok_path)

461

462 if ngrok_process.startup_error is not None:

463 raise PyngrokNgrokError("The ngrok process errored on start: {}.".format(ngrok_process.startup_error),

464 ngrok_process.logs,

465 ngrok_process.startup_error)

466 else:

467 raise PyngrokNgrokError("The ngrok process was unable to start.", ngrok_process.logs)

468

469 return ngrok_process