All code now assumes it's bases on the "ais" python module
[ais.git] / bin / inputs / udp.py
1 #!/usr/bin/python
2 # -*- coding: utf-8 -*-
3
4 '''
5 Module for receiving AIVDM data from UDP.
6 '''
7
8 import sys
9 import logging
10 import socket
11
12 from ais.inputs.common import DEFAULT_MTU, formataddr, Service, get_source_by_id4
13 from ais.inputs.stats import STATS_RATE
14 from ais.inputs.peers import udpin_get_id4_from_recvinfo
15
16 UDP_HEADER_SIZE = 28
17
18
19 class UdpChannel:
20     def __init__(self, id4, addr, port=None):
21         self.id4 = id4
22         self.source = get_source_by_id4(id4)
23         self.addr = addr
24         self.port = port
25         self.data = '' # previous line fragment that was not fully parsed
26
27     def __repr__(self):
28         result = formataddr(self.addr)
29         if self.port:
30             result += ':'+str(self.port)
31         return result
32
33 class UdpService(Service):
34
35     #TODO: Check there is not too many ports with sport_discriminate==True
36     MAX_SPORT_DISCRIMINATE = 100
37
38     def __init__(self, stdout, serverport, sport_discriminate):
39         Service.__init__(self, stdout)
40         self.serverport = serverport
41         self.sock = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
42         self.sock.settimeout(STATS_RATE)
43         self.sock.bind(('', serverport))
44         self.sport_discriminate = sport_discriminate
45         self.channels = {}
46         logging.info('Listening on UDP port %s.', serverport)
47
48
49     def _get_channel_by_addr(self, recv_from_info):
50         '''
51         Returns an element of self.channels.
52         Adds one when necessary.
53         Normally just use the IP address that is recv_from_info[0].
54         if sport_discriminate also use the source port in the hash key.
55         '''
56         addr, port, flowinfo, scopeid = recv_from_info
57         if self.sport_discriminate:
58             channelid = addr, port
59         else:
60             channelid = addr
61         if flowinfo != 0 or scopeid != 0:
62             logging.info('(DEBUG) Weird rcv_from_info: %s', recv_from_info)
63
64         channel = self.channels.get(channelid, None)
65         if channel is None:
66             id4 = udpin_get_id4_from_recvinfo(self.serverport, recv_from_info)
67             if id4 is None:
68                 logging.warning(
69                     'Rejected datagram from unauthorized source %s:%s',
70                     recv_from_info[0], recv_from_info[1])
71                 return None
72                 # TODO: Or we keep track of them low level?
73             if self.sport_discriminate:
74                 channel = UdpChannel(id4, addr, port)
75             else:
76                 channel = UdpChannel(id4, addr)
77             self.channels[channelid] = channel
78             logging.info('New connection from %s', channel)
79             
80         return channel
81
82
83     def get_activity(self):
84         while True:
85             try:
86                 data, recv_from_info = self.sock.recvfrom(DEFAULT_MTU)
87             except socket.timeout:
88                 return
89             channel = self._get_channel_by_addr(recv_from_info)
90             if channel is None:
91                 # Unauthorized
92                 return
93             channel.source.stats.npackets += 1
94             channel.source.stats.nbytes += len(data)
95             channel.source.stats.nbytes_ethernet += len(data) + UDP_HEADER_SIZE
96             from_ = '%s:%s' % (formataddr(channel.addr), recv_from_info[1])
97             logging.debug('IN %s %s', from_, repr(data))
98             channel.data += data
99             yield from_, channel
100
101
102 def main():
103     from optparse import OptionParser
104     parser = OptionParser('%prog [options] port')
105     parser.add_option('-d', '--debug',
106         help="debug mode",
107         action='store_true', dest='debug', default=False)
108     #parser.add_option('--id4',
109     #    help='4 letter string for identifying the source.',
110     #    action='store', type='str', dest='id4', default=None)
111     parser.add_option('--stdout',
112         help="Print incoming packets to stdout",
113         action='store_true', dest='stdout', default=False)
114     parser.add_option('--src-port',
115         help="Consider (ip1, port1) and (ip1, port2) different sources",
116         action='store_true', dest='sport_discriminate', default=False)
117     options, args = parser.parse_args()
118
119     if options.debug:
120         loglevel = logging.DEBUG
121     else:
122         loglevel = logging.INFO
123     logging.basicConfig(level=loglevel,
124         format='%(asctime)s %(levelname)s %(message)s')
125
126     #if options.id4:
127     #    assert len(options.id4)==4, 'ID4 must be 4 characters long.'
128     #else:
129     #    logging.warning('Source has no ID4, data will not be logged.')
130
131     if len(args)!=1:
132         print >> sys.stderr, "Missing port"
133         sys.exit(1)
134     
135     serverport = int(args[0])
136     service = UdpService(options.stdout, serverport, options.sport_discriminate)
137
138     try:
139         service.run()   
140     except KeyboardInterrupt:
141         logging.critical('Received SIGINT. Shuting down.')
142
143
144 if __name__ == '__main__':
145     main()