Coverage for pyngrok/installer.py: 94.52%

146 statements  

« prev     ^ index     » next       coverage.py v7.2.3, created at 2023-04-12 13:51 +0000

1import copy 

2import logging 

3import os 

4import platform 

5import socket 

6import sys 

7import tempfile 

8import time 

9import zipfile 

10from http import HTTPStatus 

11from urllib.request import urlopen 

12 

13import yaml 

14 

15from pyngrok.exception import PyngrokNgrokInstallError, PyngrokSecurityError, PyngrokError 

16 

17__author__ = "Alex Laird" 

18__copyright__ = "Copyright 2023, Alex Laird" 

19__version__ = "5.2.2" 

20 

21logger = logging.getLogger(__name__) 

22 

23CDN_URL_PREFIX = "https://bin.equinox.io/c/4VmDzA7iaHb/" 

24CDN_V3_URL_PREFIX = "https://bin.equinox.io/c/bNyj1mQVY4c/" 

25PLATFORMS = { 

26 "darwin_x86_64": CDN_URL_PREFIX + "ngrok-stable-darwin-amd64.zip", 

27 "darwin_x86_64_arm": CDN_URL_PREFIX + "ngrok-stable-darwin-arm64.zip", 

28 "windows_x86_64": CDN_URL_PREFIX + "ngrok-stable-windows-amd64.zip", 

29 "windows_i386": CDN_URL_PREFIX + "ngrok-stable-windows-386.zip", 

30 "linux_x86_64_arm": CDN_URL_PREFIX + "ngrok-stable-linux-arm64.zip", 

31 "linux_i386_arm": CDN_URL_PREFIX + "ngrok-stable-linux-arm.zip", 

32 "linux_i386": CDN_URL_PREFIX + "ngrok-stable-linux-386.zip", 

33 "linux_x86_64": CDN_URL_PREFIX + "ngrok-stable-linux-amd64.zip", 

34 "freebsd_x86_64": CDN_URL_PREFIX + "ngrok-stable-freebsd-amd64.zip", 

35 "freebsd_i386": CDN_URL_PREFIX + "ngrok-stable-freebsd-386.zip", 

36 "cygwin_x86_64": CDN_URL_PREFIX + "ngrok-stable-windows-amd64.zip", 

37} 

38PLATFORMS_V3 = { 

39 "darwin_x86_64": CDN_V3_URL_PREFIX + "ngrok-v3-stable-darwin-amd64.zip", 

40 "darwin_x86_64_arm": CDN_V3_URL_PREFIX + "ngrok-v3-stable-darwin-arm64.zip", 

41 "windows_x86_64": CDN_V3_URL_PREFIX + "ngrok-v3-stable-windows-amd64.zip", 

42 "windows_i386": CDN_V3_URL_PREFIX + "ngrok-v3-stable-windows-386.zip", 

43 "linux_x86_64_arm": CDN_V3_URL_PREFIX + "ngrok-v3-stable-linux-arm64.zip", 

44 "linux_i386_arm": CDN_V3_URL_PREFIX + "ngrok-v3-stable-linux-arm.zip", 

45 "linux_i386": CDN_V3_URL_PREFIX + "ngrok-v3-stable-linux-386.zip", 

46 "linux_x86_64": CDN_V3_URL_PREFIX + "ngrok-v3-stable-linux-amd64.zip", 

47 "freebsd_x86_64": CDN_V3_URL_PREFIX + "ngrok-v3-stable-freebsd-amd64.zip", 

48 "freebsd_i386": CDN_V3_URL_PREFIX + "ngrok-v3-stable-freebsd-386.zip", 

49 "cygwin_x86_64": CDN_V3_URL_PREFIX + "ngrok-v3-stable-windows-amd64.zip", 

50} 

51SUPPORTED_NGROK_VERSIONS = ["v2", "v3"] 

52DEFAULT_DOWNLOAD_TIMEOUT = 6 

53DEFAULT_RETRY_COUNT = 0 

54 

55_config_cache = {} 

56_print_progress_enabled = True 

57 

58 

59def get_ngrok_bin(): 

60 """ 

61 Get the ``ngrok`` executable for the current system. 

62 

63 :return: The name of the ``ngrok`` executable. 

64 :rtype: str 

65 """ 

66 system = platform.system().lower() 

67 if system in ["darwin", "linux", "freebsd"]: 

68 return "ngrok" 

69 elif system in ["windows", "cygwin"]: # pragma: no cover 

70 return "ngrok.exe" 

71 else: # pragma: no cover 

72 raise PyngrokNgrokInstallError("\"{}\" is not a supported platform".format(system)) 

73 

74 

75def install_ngrok(ngrok_path, ngrok_version="v2", **kwargs): 

76 """ 

77 Download and install the latest ``ngrok`` for the current system, overwriting any existing contents 

78 at the given path. 

79 

80 :param ngrok_path: The path to where the ``ngrok`` binary will be downloaded. 

81 :type ngrok_path: str 

82 :param ngrok_version: The major version of ``ngrok`` to be installed. 

83 :type ngrok_version: str, optional 

84 :param kwargs: Remaining ``kwargs`` will be passed to :func:`_download_file`. 

85 :type kwargs: dict, optional 

86 """ 

87 logger.debug( 

88 "Installing ngrok {} to {}{} ...".format(ngrok_version, ngrok_path, ", overwriting" if os.path.exists(ngrok_path) else "")) 

89 

90 ngrok_dir = os.path.dirname(ngrok_path) 

91 

92 if not os.path.exists(ngrok_dir): 

93 os.makedirs(ngrok_dir) 

94 

95 arch = "x86_64" if sys.maxsize > 2 ** 32 else "i386" 

96 if platform.uname()[4].startswith("arm") or \ 

97 platform.uname()[4].startswith("aarch64"): 

98 arch += "_arm" 

99 system = platform.system().lower() 

100 if "cygwin" in system: 

101 system = "cygwin" 

102 

103 plat = system + "_" + arch 

104 try: 

105 if ngrok_version == "v2": 

106 url = PLATFORMS[plat] 

107 elif ngrok_version == "v3": 

108 url = PLATFORMS_V3[plat] 

109 else: 

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

111 

112 logger.debug("Platform to download: {}".format(plat)) 

113 except KeyError: 

114 raise PyngrokNgrokInstallError("\"{}\" is not a supported platform".format(plat)) 

115 

116 try: 

117 download_path = _download_file(url, **kwargs) 

118 

119 _install_ngrok_zip(ngrok_path, download_path) 

120 except Exception as e: 

121 raise PyngrokNgrokInstallError("An error occurred while downloading ngrok from {}: {}".format(url, e)) 

122 

123 

124def _install_ngrok_zip(ngrok_path, zip_path): 

125 """ 

126 Extract the ``ngrok`` zip file to the given path. 

127 

128 :param ngrok_path: The path where ``ngrok`` will be installed. 

129 :type ngrok_path: str 

130 :param zip_path: The path to the ``ngrok`` zip file to be extracted. 

131 :type zip_path: str 

132 """ 

133 _print_progress("Installing ngrok ... ") 

134 

135 with zipfile.ZipFile(zip_path, "r") as zip_ref: 

136 logger.debug("Extracting ngrok binary from {} to {} ...".format(zip_path, ngrok_path)) 

137 zip_ref.extractall(os.path.dirname(ngrok_path)) 

138 

139 os.chmod(ngrok_path, int("777", 8)) 

140 

141 _clear_progress() 

142 

143 

144def get_ngrok_config(config_path, use_cache=True, ngrok_version="v2"): 

145 """ 

146 Get the ``ngrok`` config from the given path. 

147 

148 :param config_path: The ``ngrok`` config path to read. 

149 :type config_path: str 

150 :param use_cache: Use the cached version of the config (if populated). 

151 :type use_cache: bool 

152 :param ngrok_version: The major version of ``ngrok`` installed. 

153 :type ngrok_version: str, optional 

154 :return: The ``ngrok`` config. 

155 :rtype: dict 

156 """ 

157 if config_path not in _config_cache or not use_cache: 

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

159 config = yaml.safe_load(config_file) 

160 if config is None: 

161 config = get_default_config(ngrok_version) 

162 

163 _config_cache[config_path] = config 

164 

165 return _config_cache[config_path] 

166 

167 

168def get_default_config(ngrok_version): 

169 """ 

170 Get the default config params for the given major version of ``ngrok``. 

171 

172 :param ngrok_version: The major version of ``ngrok`` installed. 

173 :type ngrok_version: str, optional 

174 :return: The default config. 

175 :rtype: dict 

176 """ 

177 if ngrok_version == "v2": 

178 return {} 

179 elif ngrok_version == "v3": 

180 return {"version": "2", "region": "us"} 

181 else: 

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

183 

184 

185def install_default_config(config_path, data=None, ngrok_version="v2"): 

186 """ 

187 Install the given data to the ``ngrok`` config. If a config is not already present for the given path, create one. 

188 Before saving new data to the default config, validate that they are compatible with ``pyngrok``. 

189 

190 :param config_path: The path to where the ``ngrok`` config should be installed. 

191 :type config_path: str 

192 :param data: A dictionary of things to add to the default config. 

193 :type data: dict, optional 

194 :param ngrok_version: The major version of ``ngrok`` installed. 

195 :type ngrok_version: str, optional 

196 """ 

197 if data is None: 

198 data = {} 

199 else: 

200 data = copy.deepcopy(data) 

201 

202 data.update(get_default_config(ngrok_version)) 

203 

204 config_dir = os.path.dirname(config_path) 

205 if not os.path.exists(config_dir): 

206 os.makedirs(config_dir) 

207 if not os.path.exists(config_path): 

208 open(config_path, "w").close() 

209 

210 config = get_ngrok_config(config_path, use_cache=False) 

211 

212 config.update(data) 

213 

214 validate_config(config) 

215 

216 with open(config_path, "w") as config_file: 

217 logger.debug("Installing default ngrok config to {} ...".format(config_path)) 

218 

219 yaml.dump(config, config_file) 

220 

221 

222def validate_config(data): 

223 """ 

224 Validate that the given dict of config items are valid for ``ngrok`` and ``pyngrok``. 

225 

226 :param data: A dictionary of things to be validated as config items. 

227 :type data: dict 

228 """ 

229 if data.get("web_addr", None) is False: 

230 raise PyngrokError("\"web_addr\" cannot be False, as the ngrok API is a dependency for pyngrok") 

231 elif data.get("log_format") == "json": 

232 raise PyngrokError("\"log_format\" must be \"term\" to be compatible with pyngrok") 

233 elif data.get("log_level", "info") not in ["info", "debug"]: 

234 raise PyngrokError("\"log_level\" must be \"info\" to be compatible with pyngrok") 

235 

236 

237def _download_file(url, retries=0, **kwargs): 

238 """ 

239 Download a file to a temporary path and emit a status to stdout (if possible) as the download progresses. 

240 

241 :param url: The URL to download. 

242 :type url: str 

243 :param retries: The retry attempt index, if download fails. 

244 :type retries: int, optional 

245 :param kwargs: Remaining ``kwargs`` will be passed to :py:func:`urllib.request.urlopen`. 

246 :type kwargs: dict, optional 

247 :return: The path to the downloaded temporary file. 

248 :rtype: str 

249 """ 

250 kwargs["timeout"] = kwargs.get("timeout", DEFAULT_DOWNLOAD_TIMEOUT) 

251 

252 if not url.lower().startswith("http"): 

253 raise PyngrokSecurityError("URL must start with \"http\": {}".format(url)) 

254 

255 try: 

256 _print_progress("Downloading ngrok ...") 

257 

258 logger.debug("Download ngrok from {} ...".format(url)) 

259 

260 local_filename = url.split("/")[-1] 

261 response = urlopen(url, **kwargs) 

262 

263 status_code = response.getcode() 

264 

265 if status_code != HTTPStatus.OK: 

266 logger.debug("Response status code: {}".format(status_code)) 

267 

268 return None 

269 

270 length = response.getheader("Content-Length") 

271 if length: 

272 length = int(length) 

273 chunk_size = max(4096, length // 100) 

274 else: 

275 chunk_size = 64 * 1024 

276 

277 download_path = os.path.join(tempfile.gettempdir(), local_filename) 

278 with open(download_path, "wb") as f: 

279 size = 0 

280 while True: 

281 buffer = response.read(chunk_size) 

282 

283 if not buffer: 

284 break 

285 

286 f.write(buffer) 

287 size += len(buffer) 

288 

289 if length: 

290 percent_done = int((float(size) / float(length)) * 100) 

291 _print_progress("Downloading ngrok: {}%".format(percent_done)) 

292 

293 _clear_progress() 

294 

295 return download_path 

296 except socket.timeout as e: 

297 if retries < DEFAULT_RETRY_COUNT: 

298 logger.warning("ngrok download failed, retrying in 0.5 seconds ...") 

299 time.sleep(0.5) 

300 

301 return _download_file(url, retries + 1, **kwargs) 

302 else: 

303 raise e 

304 

305 

306def _print_progress(line): 

307 if _print_progress_enabled: 

308 sys.stdout.write("{}\r".format(line)) 

309 sys.stdout.flush() 

310 

311 

312def _clear_progress(spaces=100): 

313 if _print_progress_enabled: 

314 sys.stdout.write((" " * spaces) + "\r") 

315 sys.stdout.flush()