systematic use of context managers for dealing with files instead of open()/close...
[nepi.git] / src / nepi / resources / linux / scripts / fd-udp-connect.py
1 import base64
2 import errno
3 import os
4 import passfd
5 import signal
6 import socket
7 import time
8 import tunchannel
9
10 from optparse import OptionParser
11
12 IFF_TUN     = 0x0001
13 IFF_TAP     = 0x0002
14
15 # Trak SIGTERM, and set global termination flag instead of dying
16 TERMINATE = []
17 STARTED = False
18
19 def _finalize(sig,frame):
20     global TERMINATE
21     global STARTED
22     
23     if STARTED:
24         TERMINATE.append(None)
25     else:
26         signal.signal(signal.SIGTERM, signal.SIG_DFL)
27         os.kill(os.getpid(), signal.SIGTERM)
28
29 signal.signal(signal.SIGTERM, _finalize)
30
31 # SIGUSR1 suspends forwading, SIGUSR2 resumes forwarding
32 SUSPEND = []
33 def _suspend(sig,frame):
34     global SUSPEND
35     if not SUSPEND:
36         SUSPEND.append(None)
37 signal.signal(signal.SIGUSR1, _suspend)
38
39 def _resume(sig,frame):
40     global SUSPEND
41     if SUSPEND:
42         SUSPEND.remove(None)
43 signal.signal(signal.SIGUSR2, _resume)
44
45 def get_options():
46     usage = ("usage: %prog -a <address> -t <vif-type> -n "
47             "-b <bwlimit> -c <cipher> -k <cipher-key> "
48             "-q <txqueuelen> -p <local-port-file> -P <remote-port-file> "
49             "-o <local-ip> -O <remote-ip> -r <ret-file> ")
50     
51     parser = OptionParser(usage = usage)
52
53     parser.add_option("-a", "--address", dest="address",
54         help="Socket address to send file descriptor to", type="str")
55     parser.add_option("-t", "--vif-type", dest="vif_type",
56         help="Virtual interface type. Either IFF_TAP or IFF_TUN. "
57             "Defaults to IFF_TAP. ", default=IFF_TAP, type="str")
58     parser.add_option("-n", "--pi", dest="pi", action="store_true", 
59             default=False, help="Enable PI header")
60
61     parser.add_option("-b", "--bwlimit", dest="bwlimit",
62         help="Specifies the interface's emulated bandwidth in bytes ",
63         default=None, type="int")
64     parser.add_option("-q", "--txqueuelen", dest="txqueuelen",
65         help="Specifies the interface's transmission queue length. ",
66         default=1000, type="int")
67     parser.add_option("-c", "--cipher", dest="cipher",
68         help="Cipher to encript communication. "
69             "One of PLAIN, AES, Blowfish, DES, DES3. ",
70         default=None, type="str")
71     parser.add_option("-k", "--cipher-key", dest="cipher_key",
72         help="Specify a symmetric encryption key with which to protect "
73             "packets across the tunnel. python-crypto must be installed "
74             "on the system." ,
75         default=None, type="str")
76
77     parser.add_option("-p", "--local-port-file", dest="local_port_file",
78         help = "File where to store the local binded UDP port number ", 
79         default = "local_port_file", type="str")
80     parser.add_option("-P", "--remote-port-file", dest="remote_port_file",
81         help = "File where to read the remote UDP port number to connect to", 
82         default = "remote_port_file", type="str")
83     parser.add_option("-o", "--local-ip", dest="local_ip",
84         help = "Local host IP", type="str")
85     parser.add_option("-O", "--remote-ip", dest="remote_ip",
86         help = "Remote host IP", type="str")
87     parser.add_option("-R", "--ret-file", dest="ret_file",
88         help = "File where to store return code (success of connection) ", 
89         default = "ret_file", type="str")
90
91     (options, args) = parser.parse_args()
92             
93     vif_type = IFF_TAP
94     if options.vif_type and options.vif_type == "IFF_TUN":
95         vif_type = IFF_TUN
96   
97     address = base64.b64decode(options.address)
98
99     return (address,  vif_type, options.pi, 
100             options.local_port_file, options.remote_port_file, 
101             options.local_ip, options.remote_ip, options.ret_file, 
102             options.bwlimit, options.cipher, options.cipher_key, 
103             options.txqueuelen)
104
105 if __name__ == '__main__':
106
107     (address, vif_type, pi, local_port_file, remote_port_file, 
108             local_ip, remote_ip, ret_file, bwlimit, cipher, cipher_key, 
109             txqueuelen) = get_options()
110
111     # Create a local socket to stablish the tunnel connection
112     rsock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, 0)
113     rsock.bind((local_ip, 0))
114     (local_host, local_port) = rsock.getsockname()
115
116     # Save local port information to file
117     with open(local_port_file, 'w') as f:
118         f.write("%d\n" % local_port)
119
120     # Wait until remote port information is available
121     while not os.path.exists(remote_port_file):
122         time.sleep(2)
123
124     remote_port = ''
125     # Read remote port from file
126     # Try until something is read...
127     # xxx: There seems to be a weird behavior where
128     #       even if the file exists and had the port number,
129     #       the read operation returns empty string!
130     #       Maybe a race condition?
131     for i in xrange(10):
132         with open(remote_port_file, 'r') as f:
133             remote_port = f.read()
134
135         if remote_port:
136             break
137         
138         time.sleep(2)
139     
140     remote_port = remote_port.strip()
141     remote_port = int(remote_port)
142
143     # Connect local socket to remote port
144     rsock.connect((remote_ip, remote_port))
145     remote = os.fdopen(rsock.fileno(), 'r+b', 0)
146
147     # create local socket to pass to fd-net-device, the other is for the tunnel
148     fd, tun = socket.socketpair(socket.AF_UNIX, socket.SOCK_STREAM, 0)
149
150     # pass one end of the socket pair to the fd-net-device
151     sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
152     sock.connect(address)
153     passfd.sendfd(sock, fd, '0')
154
155     # TODO: Test connectivity!    
156
157     # Create a ret_file to indicate success
158     with open(ret_file, 'w') as f:
159         f.write("0")
160
161     STARTED = True
162
163     # Establish tunnel
164     tunchannel.tun_fwd(tun, remote,
165         with_pi = pi, 
166         ether_mode = (vif_type == IFF_TAP),
167         udp = True,
168         cipher_key = cipher_key,
169         cipher = cipher,
170         TERMINATE = TERMINATE,
171         SUSPEND = SUSPEND,
172         tunqueue = txqueuelen,
173         tunkqueue = 500,
174         bwlimit = bwlimit
175     ) 
176