7447036467947689e0d9e9ecdd50e941fb22f9ff
[ais.git] / bin / inputs / tcpout.py
1 #!/usr/bin/env python
2
3 '''
4 Module for receiving AIVDM data from outbound TCP connection.
5 '''
6
7 import logging
8 import socket
9
10 from ais.inputs.common import DEFAULT_MTU, get_source_by_id4
11 from ais.inputs.virtual import Service
12 from ais.inputs.stats import STATS_RATE
13
14 TCP_HEADER_SIZE = 52
15
16 class TcpOutChannel:
17     '''
18     Channel bound to a TCP server and listen for data there.
19     '''
20     def __init__(self, id4, hostname, port):
21         self.id4 = id4
22         self.source = get_source_by_id4(id4)
23         self.hostname = hostname
24         self.port = port
25         self.name = '%s:%s' % (self.hostname, self.port)
26         address_out = (self.hostname, self.port)
27         self.data = ''
28
29         # TODO try all IP addresses
30         # see http://docs.python.org/library/socket.html#example
31         try:
32             self.sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
33             self.sock.connect(address_out)
34             logging.info("Connected to %s using tcp6", address_out)
35         except socket.gaierror, err:
36             logging.error("Can't connet to %s: %s using tcp6", address_out, err)
37             self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
38             self.sock.connect(address_out)
39             logging.info("Connected to %s using tcp4", address_out)
40         self.sock.settimeout(STATS_RATE)
41
42     def fill_buffer(self):
43         '''
44         Grab more bytes into self.data internal buffer.
45         '''
46         try:
47             data = self.sock.recv(DEFAULT_MTU)
48         except socket.timeout:
49             return
50         if not data:
51             # FIXME: do reconnect
52             assert data, 'Connection closed'
53         stats = self.source.stats
54         stats.npackets += 1
55         stats.nbytes += len(data)
56         stats.nbytes_ethernet += len(data) + TCP_HEADER_SIZE
57         logging.debug('IN %s %s', self.name, repr(data))
58         self.data += data
59
60
61 class TcpOutService(Service):
62     '''
63     Service that connects to a TCP server and listen for data there.
64     '''
65     def __init__(self, stdout, id4, hostname, port):
66         Service.__init__(self, stdout)
67         self.channel = TcpOutChannel(id4, hostname, port)
68         
69     def fileno(self):
70         '''
71         Returns file descriptor of underlying socket, for os.select().
72         '''
73         return self.channel.sock.fileno()
74
75     def get_activity(self):
76         '''
77         Pool the channel and yields when ready.
78         '''
79         channel = self.channel
80         while True:
81             channel.fill_buffer()
82             yield channel.name, channel
83     
84
85 def main():
86     from optparse import OptionParser
87     parser = OptionParser()
88     parser.add_option('-d', '--debug',
89         help="debug mode",
90         action='store_true', dest='debug', default=False)
91     parser.add_option('--id4',
92         help='4 letter string for identifying the source.',
93         action='store', type='str', dest='id4', default=None)
94     parser.add_option('--stdout',
95         help="Print incoming packets to stdout",
96         action='store_true', dest='stdout', default=False)
97     options, args = parser.parse_args()
98     
99     if options.debug:
100         loglevel = logging.DEBUG
101     else:
102         loglevel = logging.INFO
103     logging.basicConfig(level=loglevel,
104         format='%(asctime)s %(levelname)s %(message)s')
105
106     if options.id4:
107         assert len(options.id4)==4, 'ID4 must be 4 characters long.'
108     else:
109         logging.warning('Source has no ID4, data will not be logged.')
110
111     hostname, port = args[0].split(':', 1)
112     port = int(port)
113     service = TcpOutService(options.stdout, options.id4, hostname, port)
114     try:
115         service.run()
116     except KeyboardInterrupt:
117         logging.critical('Received SIGINT. Shuting down.')
118
119 if __name__ == "__main__":
120     main()