Coverage for frappe_manager / ngrok.py: 18%

1#!/usr/bin/env python3 

2 

3import asyncio 

4import signal 

5import sys 

6import time 

7 

8import ngrok 

9 

10from frappe_manager.output_manager import get_global_output_handler, set_global_output_handler, spinner 

11from frappe_manager.output_manager.rich_output import RichOutputHandler 

12 

13 

14def create_tunnel(site_name: str, auth_token: str, port: int = 80) -> None: 

15 """ 

16 Create an ngrok HTTP tunnel for the specified site name and keep it running. 

17 

18 Args: 

19 site_name: The site name to use for host header 

20 auth_token: Ngrok authentication token 

21 port: The local port to tunnel to (default: 80) 

22 """ 

23 try: 

24 output = get_global_output_handler() 

25 except RuntimeError: 

26 output = RichOutputHandler() 

27 set_global_output_handler(output) 

28 

29 with spinner(output, f"Forwarding all requests from {site_name}"): 

30 try: 

31 ngrok.set_auth_token(auth_token) 

32 

33 listener = ngrok.forward( 

34 port=port, 

35 authtoken=auth_token, 

36 request_header_add=[f"Host: {site_name}"], 

37 opts={"addr": str(port), "host_header": site_name}, 

38 ) 

39 

40 tunnel_url = listener.url() 

41 except Exception as e: 

42 print(f"Error creating tunnel: {e}") 

43 return 

44 

45 print(f"Ingress established at: {tunnel_url}") 

46 

47 def signal_handler(sig, frame): 

48 print("\nShutting down ngrok tunnel...") 

49 listener.close() 

50 sys.exit(0) 

51 

52 signal.signal(signal.SIGINT, signal_handler) 

53 

54 try: 

55 while True: 

56 time.sleep(1) 

57 except KeyboardInterrupt: 

58 listener.close() 

59 sys.exit(0) 

60 except Exception as e: 

61 print(f"Error in tunnel: {e}") 

62 listener.close() 

63 sys.exit(1) 

64 

65 

66async def start_tunnel(site_name: str, auth_token: str): 

67 """ 

68 Start an ngrok tunnel and keep it running until interrupted. 

69 

70 Args: 

71 site_name: The site name to use for host header 

72 auth_token: Ngrok authentication token 

73 """ 

74 listener = await ngrok.connect( 

75 80, 

76 authtoken=auth_token, 

77 request_header_add=[f"Host: {site_name}"], 

78 opts={"addr": "80", "host_header": site_name}, 

79 ) 

80 

81 print(f"Ingress established at: {listener.url()}") 

82 

83 # Handle graceful shutdown 

84 def signal_handler(sig, frame): 

85 print("\nShutting down ngrok tunnel...") 

86 asyncio.create_task(listener.close()) 

87 sys.exit(0) 

88 

89 signal.signal(signal.SIGINT, signal_handler) 

90 

91 # Keep the tunnel open 

92 try: 

93 while True: 

94 await asyncio.sleep(1) 

95 except Exception as e: 

96 print(f"Error: {e}") 

97 await listener.close() 

98 

99 

100if __name__ == "__main__": 

101 if len(sys.argv) != 3: 

102 print("Usage: ngrok.py <site_name> <auth_token>") 

103 sys.exit(1) 

104 asyncio.run(start_tunnel(sys.argv[1], sys.argv[2]))