Coverage for pyngrok/process.py: 96.62%

1import atexit

2import logging

3import os

4import shlex

5import subprocess

6import threading

7import time

8from http import HTTPStatus

9from urllib.request import Request, urlopen

11import yaml

13from pyngrok import conf, installer

14from pyngrok.exception import PyngrokNgrokError, PyngrokSecurityError, PyngrokError

15from pyngrok.installer import SUPPORTED_NGROK_VERSIONS

17__author__ = "Alex Laird"

19__version__ = "6.1.2"

21logger = logging.getLogger(__name__)

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

24_current_processes = {}

27class NgrokProcess:

28 """

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

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

32 :vartype proc: subprocess.Popen

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

34 :vartype pyngrok_config: PyngrokConfig

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

36 :vartype api_url: str

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

38 :vartype logs: list[NgrokLog]

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

40 :vartype startup_error: str

41 """

43 def __init__(self, proc, pyngrok_config):

44 self.proc = proc

45 self.pyngrok_config = pyngrok_config

47 self.api_url = None

48 self.logs = []

49 self.startup_error = None

51 self._tunnel_started = False

52 self._client_connected = False

53 self._monitor_thread = None

55 def __repr__(self):

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

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

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

61 @staticmethod

62 def _line_has_error(log):

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

65 def _log_startup_line(self, line):

66 """

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

68 of the ``ngrok`` process.

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

71 :type line: str

72 :return: The parsed log.

73 :rtype: NgrokLog

74 """

75 log = self._log_line(line)

77 if log is None:

78 return

79 elif self._line_has_error(log):

80 self.startup_error = log.err

81 elif log.msg:

82 # Log ngrok startup states as they come in

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

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

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

86 self._tunnel_started = True

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

88 self._client_connected = True

90 return log

92 def _log_line(self, line):

93 """

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

95 given log line.

97 :param line: The line to be processed.

98 :type line: str

99 :return: The parsed log.

100 :rtype: NgrokLog

101 """

102 log = NgrokLog(line)

103

104 if log.line == "":

105 return None

106

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

108 self.logs.append(log)

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

110 self.logs.pop(0)

111

112 if self.pyngrok_config.log_event_callback is not None:

113 self.pyngrok_config.log_event_callback(log)

114

115 return log

116

117 def healthy(self):

118 """

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

120

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

122 :rtype: bool

123 """

124 if self.api_url is None or \

125 not self._tunnel_started or \

126 not self._client_connected:

127 return False

128

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

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

131

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

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

134 response = urlopen(request)

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

136 return False

137

138 return self.proc.poll() is None

139

140 def _monitor_process(self):

141 thread = threading.current_thread()

142

143 thread.alive = True

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

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

146

147 self._monitor_thread = None

148

149 def start_monitor_thread(self):

150 """

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

152

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

154 """

155 if self._monitor_thread is None:

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

157

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

159 self._monitor_thread.daemon = True

160 self._monitor_thread.start()

161

162 def stop_monitor_thread(self):

163 """

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

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

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

167

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

169 its logs.

170 """

171 if self._monitor_thread is not None:

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

173

174 self._monitor_thread.alive = False

175

176

177class NgrokLog:

178 """

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

180

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

182 :vartype line: str

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

184 :vartype t: str

185 :var lvl: The log's level.

186 :vartype lvl: str

187 :var msg: The log's message.

188 :vartype msg: str

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

190 :vartype err: str

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

192 :vartype addr: str

193 """

194

195 def __init__(self, line):

196 self.line = line.strip()

197 self.t = None

198 self.lvl = "NOTSET"

199 self.msg = None

200 self.err = None

201 self.addr = None

202

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

204 if "=" not in i:

205 continue

206

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

208

209 if key == "lvl":

210 if not value:

211 value = self.lvl

212

213 value = value.upper()

214 if value == "CRIT":

215 value = "CRITICAL"

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

217 value = "ERROR"

218 elif value == "WARN":

219 value = "WARNING"

220

221 if not hasattr(logging, value):

222 value = self.lvl

223

224 setattr(self, key, value)

225

226 def __repr__(self):

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

228

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

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

231 attrs.remove("line")

232

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

234

235

236def set_auth_token(pyngrok_config, token):

237 """

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

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

240

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

242 :type pyngrok_config: PyngrokConfig

243 :param token: The auth token to set.

244 :type token: str

245 """

246 if pyngrok_config.ngrok_version == "v2":

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

248 elif pyngrok_config.ngrok_version == "v3":

249 start = [pyngrok_config.ngrok_path, "config", "add-authtoken", token, "--log=stdout"]

250 else:

251 raise PyngrokError("\"ngrok_version\" must be a supported version: {}".format(SUPPORTED_NGROK_VERSIONS))

252

253 if pyngrok_config.config_path:

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

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

256 else:

257 logger.info(

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

259

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

261

262 if "Authtoken saved" not in result:

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

264

265

266def is_process_running(ngrok_path):

267 """

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

269

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

271 :type ngrok_path: str

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

273 """

274 if ngrok_path in _current_processes:

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

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

277 return True

278 else:

279 logger.debug(

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

281

282 _current_processes.pop(ngrok_path, None)

283

284 return False

285

286

287def get_process(pyngrok_config):

288 """

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

290

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

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

293

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

295 :type pyngrok_config: PyngrokConfig

296 :return: The ``ngrok`` process.

297 :rtype: NgrokProcess

298 """

299 if is_process_running(pyngrok_config.ngrok_path):

300 return _current_processes[pyngrok_config.ngrok_path]

301

302 return _start_process(pyngrok_config)

303

304

305def kill_process(ngrok_path):

306 """

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

308 issue a kill request.

309

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

311 :type ngrok_path: str

312 """

313 if is_process_running(ngrok_path):

314 ngrok_process = _current_processes[ngrok_path]

315

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

317

318 try:

319 ngrok_process.proc.kill()

320 ngrok_process.proc.wait()

321 except OSError as e: # pragma: no cover

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

323 if e.errno != 3:

324 raise e

325

326 _current_processes.pop(ngrok_path, None)

327 else:

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

329

330

331def run_process(ngrok_path, args):

332 """

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

334

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

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

337

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

339 :type ngrok_path: str

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

341 :type args: list[str]

342 """

343 _validate_path(ngrok_path)

344

345 start = [ngrok_path] + args

346 subprocess.call(start)

347

348

349def capture_run_process(ngrok_path, args):

350 """

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

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

353

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

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

356

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

358 :type ngrok_path: str

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

360 :type args: list[str]

361 :return: The output from the process.

362 :rtype: str

363 """

364 _validate_path(ngrok_path)

365

366 start = [ngrok_path] + args

367 output = subprocess.check_output(start)

368

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

370

371

372def _validate_path(ngrok_path):

373 """

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

375 relevant exception.

376

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

378 :type ngrok_path: str

379 """

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

381 raise PyngrokNgrokError(

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

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

384

385 if ngrok_path in _current_processes:

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

387

388

389def _validate_config(config_path):

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

391 config = yaml.safe_load(config_file)

392

393 if config is not None:

394 installer.validate_config(config)

395

396

397def _terminate_process(process):

398 if process is None:

399 return

400

401 try:

402 process.terminate()

403 except OSError: # pragma: no cover

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

405

406

407def _start_process(pyngrok_config):

408 """

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

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

411

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

413 :type pyngrok_config: PyngrokConfig

414 :return: The ``ngrok`` process.

415 :rtype: NgrokProcess

416 """

417 config_path = conf.get_config_path(pyngrok_config)

418

419 _validate_path(pyngrok_config.ngrok_path)

420 _validate_config(config_path)

421

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

423 if pyngrok_config.config_path:

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

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

426 if pyngrok_config.auth_token:

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

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

429 if pyngrok_config.region:

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

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

432

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

434 if os.name == "posix":

435 popen_kwargs.update(start_new_session=pyngrok_config.start_new_session)

436 elif pyngrok_config.start_new_session:

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

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

439 atexit.register(_terminate_process, proc)

440

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

442

443 ngrok_process = NgrokProcess(proc, pyngrok_config)

444 _current_processes[pyngrok_config.ngrok_path] = ngrok_process

445

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

447 while time.time() < timeout:

448 line = proc.stdout.readline()

449 ngrok_process._log_startup_line(line)

450

451 if ngrok_process.healthy():

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

453

454 ngrok_process.startup_error = None

455

456 if pyngrok_config.monitor_thread:

457 ngrok_process.start_monitor_thread()

458

459 break

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

461 break

462

463 if not ngrok_process.healthy():

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

465 kill_process(pyngrok_config.ngrok_path)

466

467 if ngrok_process.startup_error is not None:

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

469 ngrok_process.logs,

470 ngrok_process.startup_error)

471 else:

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

473

474 return ngrok_process