Coverage for hookee/tunnel.py: 91.04%

67 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-02-02 23:09 +0000

1import threading 

2import time 

3 

4from pyngrok import ngrok, conf 

5from pyngrok.conf import PyngrokConfig 

6from pyngrok.exception import PyngrokError 

7 

8__author__ = "Alex Laird" 

9__copyright__ = "Copyright 2023, Alex Laird" 

10__version__ = "2.0.0" 

11 

12 

13class Tunnel: 

14 """ 

15 An object that manages a non-blocking ``pyngrok`` tunnel and thread. 

16 

17 :var hookee_manager: Reference to the ``hookee`` Manager. 

18 :vartype hookee_manager: HookeeManager 

19 :var plugin_manager: Reference to the Plugin Manager. 

20 :vartype plugin_manager: PluginManager 

21 :var print_util: Reference to the PrintUtil. 

22 :vartype print_util: PrintUtil 

23 :var port: The server's port. 

24 :vartype port: int 

25 :var pyngrok_config: The ``pyngrok`` config. 

26 :vartype pyngrok_config: pyngrok.conf.PyngrokConfig 

27 :var public_url: The public URL of the tunnel. 

28 :vartype public_url: str 

29 :var ngrok_process: The ``ngrok`` process. 

30 :vartype ngrok_process: pyngrok.process.NgrokProcess 

31 """ 

32 

33 def __init__(self, hookee_manager): 

34 self.hookee_manager = hookee_manager 

35 self.config = self.hookee_manager.config 

36 self.plugin_manager = self.hookee_manager.plugin_manager 

37 self.print_util = self.hookee_manager.print_util 

38 self.port = self.config.get("port") 

39 

40 self.pyngrok_config = PyngrokConfig(auth_token=self.config.get("auth_token"), 

41 region=self.config.get("region")) 

42 conf.set_default(self.pyngrok_config) 

43 

44 self.public_url = None 

45 self.ngrok_process = None 

46 self._thread = None 

47 

48 def _loop(self): 

49 thread = None 

50 

51 try: 

52 self._start_tunnel() 

53 

54 thread = threading.current_thread() 

55 thread.alive = True 

56 while thread.alive: 

57 time.sleep(1) 

58 except PyngrokError: 

59 # pyngrok already logged this to the console for us 

60 pass 

61 

62 if thread: 

63 thread.alive = False 

64 

65 self.stop() 

66 

67 def start(self): 

68 """ 

69 If one is not already running, start a tunnel in a new thread. 

70 """ 

71 if self._thread is None: 

72 self.print_util.print_open_header("Opening Tunnel") 

73 

74 self._thread = threading.Thread(target=self._loop) 

75 self._thread.daemon = True 

76 self._thread.start() 

77 

78 while self.public_url is None: 

79 time.sleep(1) 

80 

81 self.print_close_header() 

82 

83 def _start_tunnel(self): 

84 options = {"bind_tls": True} 

85 subdomain = self.config.get("subdomain") 

86 hostname = self.config.get("hostname") 

87 host_header = self.config.get("host_header") 

88 auth = self.config.get("auth") 

89 if subdomain: 

90 options["subdomain"] = subdomain 

91 if hostname: 

92 options["hostname"] = hostname 

93 if host_header: 

94 options["host_header"] = host_header 

95 if auth: 

96 options["auth"] = auth 

97 

98 self.public_url = ngrok.connect(self.port, 

99 **options).public_url 

100 self.ngrok_process = ngrok.get_ngrok_process() 

101 

102 def stop(self): 

103 """ 

104 If running, kill the tunnel and cleanup its thread. 

105 """ 

106 if self._thread: 

107 ngrok.kill() 

108 

109 self.public_url = None 

110 self.ngrok_process = None 

111 self._thread = None 

112 

113 def print_close_header(self): 

114 self.print_util.print_basic( 

115 " * Tunnel: {} -> http://127.0.0.1:{}".format(self.public_url, self.port), 

116 print_when_logging=True) 

117 self.print_util.print_close_header()