Coverage for pyngrok/process.py: 94.21%

190 statements  

« prev     ^ index     » next       coverage.py v7.4.0, created at 2023-12-28 00:17 +0000

1import atexit 

2import logging 

3import os 

4import subprocess 

5import threading 

6import time 

7from http import HTTPStatus 

8from typing import Dict, List, Optional, Any 

9from urllib.request import Request, urlopen 

10 

11import yaml 

12 

13from pyngrok import conf, installer 

14from pyngrok.conf import PyngrokConfig 

15from pyngrok.exception import PyngrokNgrokError, PyngrokSecurityError, PyngrokError 

16from pyngrok.installer import SUPPORTED_NGROK_VERSIONS 

17from pyngrok.log import NgrokLog 

18 

19__author__ = "Alex Laird" 

20__copyright__ = "Copyright 2023, Alex Laird" 

21__version__ = "7.0.0" 

22 

23logger = logging.getLogger(__name__) 

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

25 

26 

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

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

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

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

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

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

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

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

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

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

332def _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] = {}