Setting the file. One moment.
Start Tunnel · Amazon Elasticache · aws/agent-toolkit-for-aws · Skills Docs
Back to the top of the page 10
Setup DevOps Agent
33
AWS Deployment
61.5
Recipe Gallery
scripts/ start_tunnel.py
Python · 344 lines · 12 KB
- An SSM-managed EC2 instance in the same VPC as the cache
17 (use find_tunnel_host.py to locate one)
18
19 Usage:
20 python start_tunnel.py \\
21 --instance-id i-0abc123def456 \\
22 --cache-host my-cache.serverless.use1.cache.amazonaws.com
23
24 python start_tunnel.py \\
25 --instance-id i-0abc123def456 \\
26 --cache-host my-cache.serverless.use1.cache.amazonaws.com \\
27 --cache-port 6379 --local-port 6379 \\
28 --region us-west-2 --profile my-profile
29 """
30
31 from __future__ import annotations
32
33 import argparse
34 import json
35 import os
36 import select
37 import shutil
38 import signal
39 import socket
40 import subprocess
41 import sys
42 import time
43
44 try :
45 import boto3
46 except ImportError :
47 print (
48 "Error: the 'boto3' package is required. \n "
49 "Install it with: \n "
50 " pip install boto3"
51 )
52 sys.exit( 2 )
53
54
55 def check_ssm_plugin () -> bool :
56 """Verify the Session Manager plugin is installed."""
57 return shutil.which( "session-manager-plugin" ) is not None
58
59
60 def check_aws_cli () -> bool :
61 """Verify the AWS CLI is installed."""
62 return shutil.which( "aws" ) is not None
63
64
65 def check_instance_ssm_status (instance_id: str , session: boto3.Session) -> bool :
66 """Verify the target instance is online in SSM."""
67 ssm = session.client( "ssm" )
68 try :
69 resp = ssm.describe_instance_information(
70 Filters = [{ "Key" : "InstanceIds" , "Values" : [instance_id]}]
71 )
72 instances = resp.get( "InstanceInformationList" , [])
73 if not instances:
74 return False
75 return instances[ 0 ].get( "PingStatus" ) == "Online"
76 except Exception as exc:
77 print ( f "Warning: could not verify SSM status: { exc } " )
78 return False
79
80
81 def check_instance_running (instance_id: str , session: boto3.Session) -> bool :
82 """Verify the target EC2 instance is in a running state."""
83 ec2 = session.client( "ec2" )
84 try :
85 resp = ec2.describe_instance_status(
86 InstanceIds = [instance_id], IncludeAllInstances = True
87 )
88 statuses = resp.get( "InstanceStatuses" , [])
89 if not statuses:
90 return False
91 return statuses[ 0 ][ "InstanceState" ][ "Name" ] == "running"
92 except Exception as exc:
93 print ( f "Warning: could not verify instance state: { exc } " )
94 return False
95
96
97 def validate_port (port: int , label: str = "port" ) -> None :
98 """Validate that a port number is in the usable range (1-65535)."""
99 if not isinstance (port, int ) or port < 1 or port > 65535 :
100 print ( f " [FAIL] { label } must be between 1 and 65535 (got { port } )" )
101 sys.exit( 1 )
102
103
104 def check_port_available (port: int ) -> bool :
105 """Check that nothing is already listening on the local port.
106
107 Returns True if the port is free, False if something is already bound.
108 """
109 try :
110 with socket.socket(socket. AF_INET , socket. SOCK_STREAM ) as s:
111 s.bind(( "127.0.0.1" , port))
112 return True
113 except OSError :
114 return False
115
116
117 def wait_for_port (port: int , timeout: float = 10.0 ) -> bool :
118 """Wait for the local port to become available (tunnel is up)."""
119 deadline = time.monotonic() + timeout
120 while time.monotonic() < deadline:
121 try :
122 with socket.create_connection(( "127.0.0.1" , port), timeout = 1.0 ):
123 return True
124 except ( ConnectionRefusedError , OSError ):
125 time.sleep( 0.5 )
126 return False
127
128
129 def start_tunnel (
130 instance_id: str ,
131 cache_host: str ,
132 cache_port: int ,
133 local_port: int ,
134 region: str ,
135 profile: str | None ,
136 ) -> None :
137 """Start SSM port forwarding and keep it running until interrupted."""
138
139 session = boto3.Session( profile_name = profile, region_name = region)
140
141 # --- Pre-flight checks ---
142 print ( "Pre-flight checks: \n " )
143
144 # 1. Port validation
145 validate_port(cache_port, "--cache-port" )
146 validate_port(local_port, "--local-port" )
147 print ( f " [OK] Port numbers are valid (local= { local_port } , remote= { cache_port } )" )
148
149 # 2. Local port not already in use
150 if not check_port_available(local_port):
151 print ( f " [FAIL] Local port { local_port } is already in use." )
152 print ( f " Another process is listening on 127.0.0.1: { local_port } ." )
153 print ( f " Stop that process first, or use --local-port to pick a different port." )
154 sys.exit( 1 )
155 print ( f " [OK] Local port { local_port } is available" )
156
157 # 3. AWS CLI
158 if not check_aws_cli():
159 print ( " [FAIL] AWS CLI not found on PATH." )
160 print ( " Install: https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html" )
161 sys.exit( 1 )
162 print ( " [OK] AWS CLI found" )
163
164 # 4. SSM plugin
165 if not check_ssm_plugin():
166 print ( " [FAIL] Session Manager plugin not found on PATH." )
167 print ( " Install: https://docs.aws.amazon.com/systems-manager/latest/userguide/session-manager-working-with-install-plugin.html" )
168 sys.exit( 1 )
169 print ( " [OK] Session Manager plugin found" )
170
171 # 5. Instance running
172 if not check_instance_running(instance_id, session):
173 print ( f " [FAIL] Instance { instance_id } is not in 'running' state." )
174 print ( " Start the instance or choose a different one (use find_tunnel_host.py)." )
175 sys.exit( 1 )
176 print ( f " [OK] Instance { instance_id } is running" )
177
178 # 6. SSM agent online
179 if not check_instance_ssm_status(instance_id, session):
180 print ( f " [WARN] Instance { instance_id } is not reporting to SSM as Online." )
181 print ( " The tunnel may fail. Ensure the instance has:" )
182 print ( " - SSM agent installed and running" )
183 print ( " - AmazonSSMManagedInstanceCore IAM policy attached" )
184 print ( " - Network route to SSM endpoints (NAT gateway or VPC endpoint)" )
185 print ( " Proceeding anyway ... \n " )
186 else :
187 print ( f " [OK] Instance { instance_id } SSM agent is online" )
188
189 print ()
190
191 # --- Build the SSM command ---
192 parameters = json.dumps(
193 {
194 "host" : [cache_host],
195 "portNumber" : [ str (cache_port)],
196 "localPortNumber" : [ str (local_port)],
197 }
198 )
199
200 cmd = [
201 "aws" , "ssm" , "start-session" ,
202 "--target" , instance_id,
203 "--document-name" , "AWS-StartPortForwardingSessionToRemoteHost" ,
204 "--parameters" , parameters,
205 "--region" , region,
206 ]
207 if profile:
208 cmd.extend([ "--profile" , profile])
209
210 print ( f "Starting tunnel: 127.0.0.1: { local_port } -> { cache_host } : { cache_port } " )
211 print ( f " via SSM instance: { instance_id } " )
212 print ( f " region: { region } " )
213 print ()
214 print ( "Tunnel command:" )
215 print ( f " { ' ' .join(cmd) } " )
216 print ()
217
218 # --- Launch the tunnel ---
219 proc = None
220 try :
221 proc = subprocess.Popen(
222 cmd,
223 stdout = subprocess. PIPE ,
224 stderr = subprocess. STDOUT ,
225 text = True ,
226 )
227
228 # Give the tunnel a moment to start, then check if the port is up
229 time.sleep( 2 )
230
231 if proc.poll() is not None :
232 # Process already exited -- something went wrong
233 output = proc.stdout.read() if proc.stdout else ""
234 print ( f "Tunnel process exited immediately (code { proc.returncode } ): \n " )
235 print (output)
236 sys.exit( 1 )
237
238 if wait_for_port(local_port, timeout = 15.0 ):
239 print ( f "Tunnel is UP. Local port 127.0.0.1: { local_port } is accepting connections. \n " )
240 is_serverless = ".serverless." in cache_host
241 if is_serverless:
242 print ( "Note: ElastiCache Serverless requires TLS. Your client must connect" )
243 print ( f " with TLS enabled and SNI set to { cache_host } " )
244 print ( f " (the tunnel forwards raw TCP; TLS negotiation happens end-to-end). \n " )
245 print ( "Test with:" )
246 print ( f " python scripts/test_connection.py 127.0.0.1 --port { local_port } \\ " )
247 print ( f " --tunnel-mode --server-name { cache_host } \\ " )
248 print ( f " --username <user> --password <password>" )
249 print ()
250 print ( "Press Ctrl+C to stop the tunnel." )
251 else :
252 print ( f "Warning: local port { local_port } not yet responding after 15s." )
253 print ( "The tunnel may still be starting. Check the output below. \n " )
254
255 # Stream SSM output with periodic health checks.
256 # Uses select() to avoid blocking forever if the session stalls.
257 health_check_interval = 30 # seconds between port probes
258
259 assert proc.stdout is not None
260 while proc.poll() is None :
261 ready, _, _ = select.select([proc.stdout], [], [], health_check_interval)
262 if ready:
263 line = proc.stdout.readline()
264 if line:
265 print ( f " [SSM] { line } " , end = "" )
266 else :
267 break # EOF -- process closed stdout
268 else :
269 # No output for health_check_interval seconds, check tunnel
270 if not wait_for_port(local_port, timeout = 3.0 ):
271 print ( f " \n [WARN] Tunnel port { local_port } is no longer responding." )
272 print ( " The SSM session may have stalled. Press Ctrl+C to stop." )
273
274 proc.wait()
275 if proc.returncode != 0 :
276 print ( f " \n Tunnel exited with code { proc.returncode } " )
277 sys.exit( 1 )
278
279 except KeyboardInterrupt :
280 print ( " \n\n Shutting down tunnel ..." )
281 if proc and proc.poll() is None :
282 proc.send_signal(signal. SIGTERM )
283 try :
284 proc.wait( timeout = 5 )
285 except subprocess.TimeoutExpired:
286 proc.kill()
287 print ( "Tunnel closed." )
288 except Exception as exc:
289 print ( f " \n Error: { exc } " )
290 if proc and proc.poll() is None :
291 proc.kill()
292 sys.exit( 1 )
293
294
295 def main () -> None :
296 parser = argparse.ArgumentParser(
297 description = "Start an SSM port-forwarding tunnel to an ElastiCache endpoint." ,
298 formatter_class = argparse.RawDescriptionHelpFormatter,
299 epilog = __doc__ ,
300 )
301 parser.add_argument(
302 "--instance-id" , required = True ,
303 help = "SSM-managed EC2 instance ID (use find_tunnel_host.py to discover)" ,
304 )
305 parser.add_argument(
306 "--cache-host" , required = True ,
307 help = "ElastiCache endpoint hostname (e.g. my-cache.serverless.use1.cache.amazonaws.com)" ,
308 )
309 parser.add_argument(
310 "--cache-port" , type = int , default = 6379 ,
311 help = "ElastiCache port (default: 6379; serverless reader endpoint uses 6380)" ,
312 )
313 parser.add_argument(
314 "--local-port" , type = int , default = 6379 ,
315 help = "Local port to forward to (default: 6379)" ,
316 )
317 _default_region = (
318 os.environ.get( "AWS_REGION" )
319 or os.environ.get( "AWS_DEFAULT_REGION" )
320 or "us-east-1"
321 )
322 parser.add_argument(
323 "--region" , default = _default_region,
324 help = f "AWS region (default: { _default_region } )" ,
325 )
326 parser.add_argument(
327 "--profile" , default = None ,
328 help = "AWS profile name" ,
329 )
330
331 args = parser.parse_args()
332
333 start_tunnel(
334 instance_id = args.instance_id,
335 cache_host = args.cache_host,
336 cache_port = args.cache_port,
337 local_port = args.local_port,
338 region = args.region,
339 profile = args.profile,
340 )
341
342
343 if __name__ == "__main__" :
344 main()