Coverage for pyngrok/installer.py: 94.52%

1 import copy

2 import logging

3 import os

4 import platform

5 import socket

6 import sys

7 import tempfile

8 import time

9 import zipfile

10 from http import HTTPStatus

11 from urllib.request import urlopen

12

13 import yaml

14

15 from pyngrok.exception import PyngrokNgrokInstallError, PyngrokSecurityError, PyngrokError

16

17 __author__ = "Alex Laird"

18 __copyright__ = "Copyright 2023, Alex Laird"

19 __version__ = "5.2.2"

20

21 logger = logging.getLogger(__name__)

22

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

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

25 PLATFORMS = {

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}

38 PLATFORMS_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}

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

52 DEFAULT_DOWNLOAD_TIMEOUT = 6

53 DEFAULT_RETRY_COUNT = 0

54

55 _config_cache = {}

56 _print_progress_enabled = True

57

58

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

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

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

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

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

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

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

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

306 def _print_progress(line):

307 if _print_progress_enabled:

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

309 sys.stdout.flush()

310

311

312 def _clear_progress(spaces=100):

313 if _print_progress_enabled:

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

315 sys.stdout.flush()