beacon.py

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