Coverage for pyngrok/installer.py: 95.31%

# Coverage for pyngrok/installer.py : 95.31%

## 128 statements

Hot-keys on this page

r m x p toggle line displays

j k next/prev highlighted chunk

0 (zero) top of page

1 (one) first highlighted chunk

1 import logging

2 import os

3 import platform

4 import socket

5 import sys

6 import tempfile

7 import time

8 import zipfile

9 from http import HTTPStatus

10 from urllib.request import urlopen

11

12 import yaml

13

14 from pyngrok.exception import PyngrokNgrokInstallError, PyngrokSecurityError, PyngrokError

15

16 __author__ = "Alex Laird"

17 __copyright__ = "Copyright 2020, Alex Laird"

18 __version__ = "5.0.0"

19

20 logger = logging.getLogger(__name__)

21

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

23 PLATFORMS = {

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

25 "darwin_i386": CDN_URL_PREFIX + "ngrok-stable-darwin-386.zip",

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

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

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

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

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

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

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

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

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

35}

36 DEFAULT_DOWNLOAD_TIMEOUT = 6

37 DEFAULT_RETRY_COUNT = 0

38

39 _config_cache = None

40 _print_progress_enabled = True

41

42

43 def get_ngrok_bin():

44 """

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

46

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

48 :rtype: str

49 """

50 system = platform.system().lower()

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

52 return "ngrok"

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

54 return "ngrok.exe"

55 else: # pragma: no cover

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

57

58

59 def install_ngrok(ngrok_path, **kwargs):

60 """

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

62 at the given path.

63

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

65 :type ngrok_path: str

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

67 :type kwargs: dict, optional

68 """

69 logger.debug(

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

71

72 ngrok_dir = os.path.dirname(ngrok_path)

73

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

75 os.makedirs(ngrok_dir)

76

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

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

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

80 arch += "_arm"

81 system = platform.system().lower()

82 if "cygwin" in system:

83 system = "cygwin"

84

85 plat = system + "_" + arch

86 try:

87 url = PLATFORMS[plat]

88

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

90 except KeyError:

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

92

93 try:

94 download_path = _download_file(url, **kwargs)

95

96 _install_ngrok_zip(ngrok_path, download_path)

97 except Exception as e:

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

99

100

101 def _install_ngrok_zip(ngrok_path, zip_path):

102 """

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

104

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

106 :type ngrok_path: str

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

108 :type zip_path: str

109 """

110 _print_progress("Installing ngrok ... ")

111

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

113 logger.debug("Extracting ngrok binary to {} ...".format(zip_path))

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

115

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

117

118 _clear_progress()

119

120

121 def get_ngrok_config(config_path, use_cache=True):

122 """

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

124

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

126 :type config_path: str

127 :param use_cache: Use the cached version of the cache (if populated).

128 :type use_cache: bool

129 :return: The ``ngrok`` config.

130 :rtype: dict

131 """

132 global _config_cache

133

134 if not _config_cache or not use_cache:

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

136 config = yaml.safe_load(config_file)

137 if config is None:

138 config = {}

139

140 _config_cache = config

141

142 return _config_cache

143

144

145 def install_default_config(config_path, data=None):

146 """

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

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

149

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

151 :type config_path: str

152 :param data: A dictionary of things to added to the default config.

153 :type data: dict, optional

154 """

155 if data is None:

156 data = {}

157

158 config_dir = os.path.dirname(config_path)

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

160 os.makedirs(config_dir)

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

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

163

164 config = get_ngrok_config(config_path, use_cache=False)

165

166 config.update(data)

167

168 validate_config(config)

169

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

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

172

173 yaml.dump(config, config_file)

174

175

176 def validate_config(data):

177 """

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

179

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

181 :type data: dict

182 """

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

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

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

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

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

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

189

190

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

192 """

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

194

195 :param url: The URL to download.

196 :type url: str

197 :param retries: The number of retries to attempt, if download fails.

198 :type retries: int, optional

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

200 :type kwargs: dict, optional

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

202 :rtype: str

203 """

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

205

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

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

208

209 try:

210 _print_progress("Downloading ngrok ...")

211

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

213

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

215 response = urlopen(url, **kwargs)

216

217 status_code = response.getcode()

218

219 if status_code != HTTPStatus.OK:

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

221

222 return None

223

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

225 if length:

226 length = int(length)

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

228 else:

229 chunk_size = 64 * 1024

230

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

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

233 size = 0

234 while True:

235 buffer = response.read(chunk_size)

236

237 if not buffer:

238 break

239

240 f.write(buffer)

241 size += len(buffer)

242

243 if length:

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

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

246

247 _clear_progress()

248

249 return download_path

250 except socket.timeout as e:

251 if retries < DEFAULT_RETRY_COUNT:

252 time.sleep(0.5)

253

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

255 else:

256 raise e

257

258

259 def _print_progress(line):

260 if _print_progress_enabled:

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

262 sys.stdout.flush()

263

264

265 def _clear_progress(spaces=100):

266 if _print_progress_enabled:

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

268 sys.stdout.flush()