[bug] ship newer client discover script
[puppet_squid_deb_proxy.git] / files / client / apt-avahi-discover
1 #!/usr/bin/python
2 #
3 # use avahi to find a _apt_proxy._tcp provider and return
4 # a http proxy string suitable for apt
5
6 import asyncore
7 import functools
8 import os
9 import socket
10 import sys
11 import time
12 from subprocess import Popen, PIPE, call
13 from syslog import syslog, LOG_INFO, LOG_USER
14
15 DEFAULT_CONNECT_TIMEOUT_SEC = 2
16
17 def DEBUG(msg):
18     if "--debug" in sys.argv:
19         sys.stderr.write(msg + "\n")
20
21
22 def get_avahi_discover_timeout():
23     APT_AVAHI_TIMEOUT_VAR = "APT::Avahi-Discover::Timeout"
24     p = Popen(
25         ["/usr/bin/apt-config", "shell", "TIMEOUT", APT_AVAHI_TIMEOUT_VAR], 
26         stdout=PIPE)
27     stdout, stderr = p.communicate()
28     if not stdout:
29         DEBUG(
30             "no timeout set, using default '%s'" % DEFAULT_CONNECT_TIMEOUT_SEC)
31         return DEFAULT_CONNECT_TIMEOUT_SEC
32     if not stdout.startswith("TIMEOUT="):
33         raise ValueError("got unexpected apt-config output: '%s'" % stdout)
34     varname, sep, value = stdout.strip().partition("=")
35     timeout = int(value.strip("'"))
36     DEBUG("using timeout: '%s'" % timeout)
37     return timeout
38
39 @functools.total_ordering
40 class AptAvahiClient(asyncore.dispatcher):
41     def __init__(self, addr):
42         asyncore.dispatcher.__init__(self)
43         if is_ipv6(addr[0]):
44             self.create_socket(socket.AF_INET6, socket.SOCK_STREAM)
45             self.connect( (addr[0], addr[1], 0, 0) )
46         else:
47             self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
48             self.connect(addr)
49         self._time_init = time.time()
50         self.time_to_connect = sys.maxint
51         self.address = addr
52     def handle_connect(self):
53         self.time_to_connect = time.time() - self._time_init
54         self.close()
55     def __eq__(self, other):
56         return self.time_to_connect == other.time_to_connect
57     def __lt__(self, other):
58         return self.time_to_connect < other.time_to_connect
59     def __repr__(self):
60         return "<%s> %s: %s" % (
61             self.__class__.__name__, self.addr, self.time_to_connect)
62     def log(self, message):
63         syslog((LOG_INFO|LOG_USER), '%s\n' % str(message))
64     def log_info(self, message, type='info'):
65         if type not in self.ignore_log_types:
66             self.log('%s: %s' % (type, message))
67
68
69 def is_ipv6(a):
70     return ':' in a
71
72 def is_linklocal(addr):
73     # Link-local should start with fe80 and six null bytes
74     return addr.startswith("fe80::")
75
76 def get_proxy_host_port_from_avahi():
77     service = '_apt_proxy._tcp'
78
79     # Obtain all of the services addresses from avahi, pulling the IPv6
80     # addresses to the top.
81     addr4 = []
82     addr6 = []
83     p = Popen(['avahi-browse', '-kprtf', service], stdout=PIPE)
84     DEBUG("avahi-browse output:")
85     for line in p.stdout:
86         DEBUG(" '%s'" % line)
87         if line.startswith('='):
88             tokens = line.split(';')
89             addr = tokens[7]
90             port = int(tokens[8])
91             if is_ipv6(addr):
92                 # We need to skip ipv6 link-local addresses since 
93                 # APT can't use them
94                 if not is_linklocal(addr):
95                     addr6.append((addr, port))
96             else:
97                 addr4.append((addr, port))
98
99     # Run through the offered addresses and see if we we have a bound local
100     # address for it.
101     addrs = []
102     for (ip, port) in addr6 + addr4:
103         try:
104             res = socket.getaddrinfo(ip, port, 0, 0, 0, socket.AI_ADDRCONFIG)
105             if res:
106                 addrs.append((ip, port))
107         except socket.gaierror:
108             pass
109     if not addrs:
110         return None
111     
112     # sort by answering speed
113     hosts = []
114     for addr in addrs:
115         hosts.append(AptAvahiClient(addr))
116     # 2s timeout, arbitray
117     timeout = get_avahi_discover_timeout()
118     asyncore.loop(timeout=timeout)
119     DEBUG("sorted hosts: '%s'" % sorted(hosts))
120
121     # No host wanted to connect
122     if (all(h.time_to_connect == sys.maxint for h in hosts)):
123         return None
124
125     fastest_host = sorted(hosts)[0]
126     fastest_address = fastest_host.address
127     return fastest_address
128
129
130 if __name__ == "__main__":
131     # Dump the approved address out in an appropriate format.
132     address = get_proxy_host_port_from_avahi()
133     if address:
134         (ip, port) = address
135         if is_ipv6(ip):
136             print "http://[%s]:%s/" % (ip, port)
137         else:
138             print "http://%s:%s/" % (ip, port)