beacon.py

   1 import logging
   2 import re
   3 import socket
   4 import struct
   5 import time
   6 from threading import Timer
   7 from urllib import quote
   8
   9 import Zeroconf
  10
  11 import config
  12 from plugin import GetPlugin
  13
  14 SHARE_TEMPLATE = '/TiVoConnect?Command=QueryContainer&Container=%s'
  15 PLATFORM_MAIN = 'pyTivo'
  16 PLATFORM_VIDEO = 'pc/pyTivo'    # For the nice icon
  17
  18 class ZCListener:
  19     def __init__(self, names):
  20         self.names = names
  21
  22     def removeService(self, server, type, name):
  23         self.names.remove(name.replace('.' + type, ''))
  24
  25     def addService(self, server, type, name):
  26         self.names.append(name.replace('.' + type, ''))
  27
  28 class ZCBroadcast:
  29     def __init__(self, logger):
  30         """ Announce our shares via Zeroconf. """
  31         self.share_names = []
  32         self.share_info = []
  33         self.logger = logger
  34         self.rz = Zeroconf.Zeroconf()
  35         self.renamed = {}
  36         old_titles = self.scan()
  37         address = socket.inet_aton(config.get_ip())
  38         port = int(config.getPort())
  39         logger.info('Announcing shares...')
  40         for section, settings in config.getShares():
  41             try:
  42                 ct = GetPlugin(settings['type']).CONTENT_TYPE
  43             except:
  44                 continue
  45             if ct.startswith('x-container/'):
  46                 if 'video' in ct:
  47                     platform = PLATFORM_VIDEO
  48                 else:
  49                     platform = PLATFORM_MAIN
  50                 logger.info('Registering: %s' % section)
  51                 self.share_names.append(section)
  52                 desc = {'path': SHARE_TEMPLATE % quote(section),
  53                         'platform': platform, 'protocol': 'http'}
  54                 tt = ct.split('/')[1]
  55                 title = section
  56                 count = 1
  57                 while title in old_titles:
  58                     count += 1
  59                     title = '%s [%d]' % (section, count)
  60                     self.renamed[section] = title
  61                 info = Zeroconf.ServiceInfo('_%s._tcp.local.' % tt,
  62                     '%s._%s._tcp.local.' % (title, tt),
  63                     address, port, 0, 0, desc)
  64                 self.rz.registerService(info)
  65                 self.share_info.append(info)
  66
  67     def scan(self):
  68         """ Look for TiVos using Zeroconf. """
  69         VIDS = '_tivo-videos._tcp.local.'
  70         names = []
  71
  72         self.logger.info('Scanning for TiVos...')
  73
  74         # Get the names of servers offering TiVo videos
  75         browser = Zeroconf.ServiceBrowser(self.rz, VIDS, ZCListener(names))
  76
  77         # Give them a second to respond
  78         time.sleep(1)
  79
  80         # Any results?
  81         if names:
  82             config.tivos_found = True
  83
  84         # Now get the addresses -- this is the slow part
  85         for name in names:
  86             info = self.rz.getServiceInfo(VIDS, name + '.' + VIDS)
  87             if info and 'TSN' in info.properties:
  88                 tsn = info.properties['TSN']
  89                 address = socket.inet_ntoa(info.getAddress())
  90                 port = info.getPort()
  91                 config.tivos[tsn] = {'name': name, 'address': address,
  92                                      'port': port}
  93                 config.tivos[tsn].update(info.properties)
  94                 self.logger.info(name)
  95
  96         return names
  97
  98     def shutdown(self):
  99         self.logger.info('Unregistering: %s' % ' '.join(self.share_names))
 100         for info in self.share_info:
 101             self.rz.unregisterService(info)
 102         self.rz.close()
 103
 104 class Beacon:
 105     def __init__(self):
 106         self.UDPSock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 107         self.UDPSock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
 108         self.services = []
 109
 110         self.platform = PLATFORM_VIDEO
 111         for section, settings in config.getShares():
 112             try:
 113                 ct = GetPlugin(settings['type']).CONTENT_TYPE
 114             except:
 115                 continue
 116             if ct in ('x-container/tivo-music', 'x-container/tivo-photos'):
 117                 self.platform = PLATFORM_MAIN
 118                 break
 119
 120         if config.get_zc():
 121             logger = logging.getLogger('pyTivo.beacon')
 122             try:
 123                 self.bd = ZCBroadcast(logger)
 124             except:
 125                 logger.error('Zeroconf failure')
 126                 self.bd = None
 127         else:
 128             self.bd = None
 129
 130     def add_service(self, service):
 131         self.services.append(service)
 132         self.send_beacon()
 133
 134     def format_services(self):
 135         return ';'.join(self.services)
 136
 137     def format_beacon(self, conntype, services=True):
 138         beacon = ['tivoconnect=1',
 139                   'method=%s' % conntype,
 140                   'identity={%s}' % config.getGUID(),
 141                   'machine=%s' % socket.gethostname(),
 142                   'platform=%s' % self.platform]
 143
 144         if services:
 145             beacon.append('services=' + self.format_services())
 146         else:
 147             beacon.append('services=TiVoMediaServer:0/http')
 148
 149         return '\n'.join(beacon) + '\n'
 150
 151     def send_beacon(self):
 152         beacon_ips = config.getBeaconAddresses()
 153         beacon = self.format_beacon('broadcast')
 154         for beacon_ip in beacon_ips.split():
 155             if beacon_ip != 'listen':
 156                 try:
 157                     packet = beacon
 158                     while packet:
 159                         result = self.UDPSock.sendto(packet, (beacon_ip, 2190))
 160                         if result < 0:
 161                             break
 162                         packet = packet[result:]
 163                 except Exception, e:
 164                     print e
 165
 166     def start(self):
 167         self.send_beacon()
 168         self.timer = Timer(60, self.start)
 169         self.timer.start()
 170
 171     def stop(self):
 172         self.timer.cancel()
 173         if self.bd:
 174             self.bd.shutdown()
 175
 176     def recv_bytes(self, sock, length):
 177         block = ''
 178         while len(block) < length:
 179             add = sock.recv(length - len(block))
 180             if not add:
 181                 break
 182             block += add
 183         return block
 184
 185     def recv_packet(self, sock):
 186         length = struct.unpack('!I', self.recv_bytes(sock, 4))[0]
 187         return self.recv_bytes(sock, length)
 188
 189     def send_packet(self, sock, packet):
 190         sock.sendall(struct.pack('!I', len(packet)) + packet)
 191
 192     def listen(self):
 193         """ For the direct-connect, TCP-style beacon """
 194         import thread
 195
 196         def server():
 197             TCPSock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 198             TCPSock.bind(('', 2190))
 199             TCPSock.listen(5)
 200
 201             while True:
 202                 # Wait for a connection
 203                 client, address = TCPSock.accept()
 204
 205                 # Accept (and discard) the client's beacon
 206                 self.recv_packet(client)
 207
 208                 # Send ours
 209                 self.send_packet(client, self.format_beacon('connected'))
 210
 211                 client.close()
 212
 213         thread.start_new_thread(server, ())
 214
 215     def get_name(self, address):
 216         """ Exchange beacons, and extract the machine name. """
 217         our_beacon = self.format_beacon('connected', False)
 218         machine_name = re.compile('machine=(.*)\n').search
 219
 220         try:
 221             tsock = socket.socket()
 222             tsock.connect((address, 2190))
 223             self.send_packet(tsock, our_beacon)
 224             tivo_beacon = self.recv_packet(tsock)
 225             tsock.close()
 226             name = machine_name(tivo_beacon).groups()[0]
 227         except:
 228             name = address
 229
 230         return name