Coverage for pyngrok/process.py: 96.17%

209 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-11-29 17:05 +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 

10 

11import yaml 

12 

13from pyngrok import conf, installer 

14from pyngrok.exception import PyngrokNgrokError, PyngrokSecurityError, PyngrokError 

15 

16__author__ = "Alex Laird" 

17__copyright__ = "Copyright 2022, Alex Laird" 

18__version__ = "5.2.0" 

19 

20from pyngrok.installer import SUPPORTED_NGROK_VERSIONS 

21 

22logger = logging.getLogger(__name__) 

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

24 

25_current_processes = {} 

26 

27 

28class 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 

178class 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 

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

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 

267def 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 

288def 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 

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. 

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 

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

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. 

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 

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. 

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 

390def _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 

398def _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 

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. 

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