Coverage for pyngrok/installer.py: 94.56%

147 statements  

« prev     ^ index     » next       coverage.py v7.3.1, created at 2023-09-20 14:14 +0000

1import copy 

2import logging 

3import os 

4import platform 

5import socket 

6import sys 

7import tempfile 

8import time 

9import zipfile 

10from http import HTTPStatus 

11from typing import Any, Dict, Optional 

12from urllib.request import urlopen 

13 

14import yaml 

15 

16from pyngrok.exception import PyngrokNgrokInstallError, PyngrokSecurityError, PyngrokError 

17 

18__author__ = "Alex Laird" 

19__copyright__ = "Copyright 2023, Alex Laird" 

20__version__ = "7.0.0" 

21 

22logger = logging.getLogger(__name__) 

23 

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

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

26PLATFORMS = { 

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

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

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

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

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

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

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

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

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

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

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

38} 

39PLATFORMS_V3 = { 

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

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

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

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

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

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

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

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

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

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

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

51} 

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

53DEFAULT_DOWNLOAD_TIMEOUT = 6 

54DEFAULT_RETRY_COUNT = 0 

55 

56_config_cache: Dict[str, Dict[str, Any]] = {} 

57_print_progress_enabled = True 

58 

59 

60def get_ngrok_bin() -> str: 

61 """ 

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

63 

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

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: str, 

76 ngrok_version: Optional[str] = "v3", 

77 **kwargs: Any) -> None: 

78 """ 

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

80 at the given path. 

81 

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

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

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

85 """ 

86 logger.debug( 

87 "Installing ngrok {} to {}{} ...".format(ngrok_version, ngrok_path, 

88 ", 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: str, 

125 zip_path: str) -> None: 

126 """ 

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

128 

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

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

131 """ 

132 _print_progress("Installing ngrok ... ") 

133 

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

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

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

137 

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

139 

140 _clear_progress() 

141 

142 

143def get_ngrok_config(config_path: str, 

144 use_cache: bool = True, 

145 ngrok_version: Optional[str] = "v3") -> Dict[str, Any]: 

146 """ 

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

148 

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

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

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

152 :return: The ``ngrok`` config. 

153 """ 

154 if config_path not in _config_cache or not use_cache: 

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

156 config = yaml.safe_load(config_file) 

157 if config is None: 

158 config = get_default_config(ngrok_version) 

159 

160 _config_cache[config_path] = config 

161 

162 return _config_cache[config_path] 

163 

164 

165def get_default_config(ngrok_version: Optional[str]) -> Dict[str, Any]: 

166 """ 

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

168 

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

170 :return: The default config. 

171 """ 

172 if ngrok_version == "v2": 

173 return {} 

174 elif ngrok_version == "v3": 

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

176 else: 

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

178 

179 

180def install_default_config(config_path: str, 

181 data: Optional[Dict[str, Any]] = None, 

182 ngrok_version: Optional[str] = "v3") -> None: 

183 """ 

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

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

186 

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

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

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

190 """ 

191 if data is None: 

192 data = {} 

193 else: 

194 data = copy.deepcopy(data) 

195 

196 data.update(get_default_config(ngrok_version)) 

197 

198 config_dir = os.path.dirname(config_path) 

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

200 os.makedirs(config_dir) 

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

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

203 

204 config = get_ngrok_config(config_path, use_cache=False, ngrok_version=ngrok_version) 

205 

206 config.update(data) 

207 

208 validate_config(config) 

209 

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

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

212 

213 yaml.dump(config, config_file) 

214 

215 

216def validate_config(data: Dict[str, Any]) -> None: 

217 """ 

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

219 

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

221 """ 

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

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

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

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

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

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

228 

229 

230def _download_file(url: str, 

231 retries: int = 0, 

232 **kwargs: Any) -> str: 

233 """ 

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

235 

236 :param url: The URL to download. 

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

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

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

240 """ 

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

242 

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

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

245 

246 try: 

247 _print_progress("Downloading ngrok ...") 

248 

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

250 

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

252 response = urlopen(url, **kwargs) 

253 

254 status_code = response.getcode() 

255 

256 if status_code != HTTPStatus.OK: 

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

258 

259 raise PyngrokNgrokInstallError("Download failed, status code: {}".format(status_code)) 

260 

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

262 if length: 

263 length = int(length) 

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

265 else: 

266 chunk_size = 64 * 1024 

267 

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

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

270 size = 0 

271 while True: 

272 buffer = response.read(chunk_size) 

273 

274 if not buffer: 

275 break 

276 

277 f.write(buffer) 

278 size += len(buffer) 

279 

280 if length: 

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

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

283 

284 _clear_progress() 

285 

286 return download_path 

287 except socket.timeout as e: 

288 if retries < DEFAULT_RETRY_COUNT: 

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

290 time.sleep(0.5) 

291 

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

293 else: 

294 raise e 

295 

296 

297def _print_progress(line: str) -> None: 

298 if _print_progress_enabled: 

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

300 sys.stdout.flush() 

301 

302 

303def _clear_progress(spaces: int = 100) -> None: 

304 if _print_progress_enabled: 

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

306 sys.stdout.flush()