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:
  90                 tsn = info.properties.get('TSN')
  91                 if config.get_server('togo_all'):
  92                     tsn = info.properties.get('tsn', tsn)
  93                 if tsn:
  94                     address = socket.inet_ntoa(info.getAddress())
  95                     port = info.getPort()
  96                     config.tivos[tsn] = {'name': name, 'address': address,
  97                                          'port': port}
  98                     config.tivos[tsn].update(info.properties)
  99                     self.logger.info(name)
 100
 101         return names
 102
 103     def shutdown(self):
 104         self.logger.info('Unregistering: %s' % ' '.join(self.share_names))
 105         for info in self.share_info:
 106             self.rz.unregisterService(info)
 107         self.rz.close()
 108
 109 class Beacon:
 110     def __init__(self):
 111         self.UDPSock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 112         self.UDPSock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
 113         self.services = []
 114
 115         self.platform = PLATFORM_VIDEO
 116         for section, settings in config.getShares():
 117             try:
 118                 ct = GetPlugin(settings['type']).CONTENT_TYPE
 119             except:
 120                 continue
 121             if ct in ('x-container/tivo-music', 'x-container/tivo-photos'):
 122                 self.platform = PLATFORM_MAIN
 123                 break
 124
 125         if config.get_zc():
 126             logger = logging.getLogger('pyTivo.beacon')
 127             try:
 128                 self.bd = ZCBroadcast(logger)
 129             except:
 130                 logger.error('Zeroconf failure')
 131                 self.bd = None
 132         else:
 133             self.bd = None
 134
 135     def add_service(self, service):
 136         self.services.append(service)
 137         self.send_beacon()
 138
 139     def format_services(self):
 140         return ';'.join(self.services)
 141
 142     def format_beacon(self, conntype, services=True):
 143         beacon = ['tivoconnect=1',
 144                   'method=%s' % conntype,
 145                   'identity={%s}' % config.getGUID(),
 146                   'machine=%s' % socket.gethostname(),
 147                   'platform=%s' % self.platform]
 148
 149         if services:
 150             beacon.append('services=' + self.format_services())
 151         else:
 152             beacon.append('services=TiVoMediaServer:0/http')
 153
 154         return '\n'.join(beacon) + '\n'
 155
 156     def send_beacon(self):
 157         beacon_ips = config.getBeaconAddresses()
 158         beacon = self.format_beacon('broadcast')
 159         for beacon_ip in beacon_ips.split():
 160             if beacon_ip != 'listen':
 161                 try:
 162                     packet = beacon
 163                     while packet:
 164                         result = self.UDPSock.sendto(packet, (beacon_ip, 2190))
 165                         if result < 0:
 166                             break
 167                         packet = packet[result:]
 168                 except Exception, e:
 169                     print e
 170
 171     def start(self):
 172         self.send_beacon()
 173         self.timer = Timer(60, self.start)
 174         self.timer.start()
 175
 176     def stop(self):
 177         self.timer.cancel()
 178         if self.bd:
 179             self.bd.shutdown()
 180
 181     def recv_bytes(self, sock, length):
 182         block = ''
 183         while len(block) < length:
 184             add = sock.recv(length - len(block))
 185             if not add:
 186                 break
 187             block += add
 188         return block
 189
 190     def recv_packet(self, sock):
 191         length = struct.unpack('!I', self.recv_bytes(sock, 4))[0]
 192         return self.recv_bytes(sock, length)
 193
 194     def send_packet(self, sock, packet):
 195         sock.sendall(struct.pack('!I', len(packet)) + packet)
 196
 197     def listen(self):
 198         """ For the direct-connect, TCP-style beacon """
 199         import thread
 200
 201         def server():
 202             TCPSock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 203             TCPSock.bind(('', 2190))
 204             TCPSock.listen(5)
 205
 206             while True:
 207                 # Wait for a connection
 208                 client, address = TCPSock.accept()
 209
 210                 # Accept (and discard) the client's beacon
 211                 self.recv_packet(client)
 212
 213                 # Send ours
 214                 self.send_packet(client, self.format_beacon('connected'))
 215
 216                 client.close()
 217
 218         thread.start_new_thread(server, ())
 219
 220     def get_name(self, address):
 221         """ Exchange beacons, and extract the machine name. """
 222         our_beacon = self.format_beacon('connected', False)
 223         machine_name = re.compile('machine=(.*)\n').search
 224
 225         try:
 226             tsock = socket.socket()
 227             tsock.connect((address, 2190))
 228             self.send_packet(tsock, our_beacon)
 229             tivo_beacon = self.recv_packet(tsock)
 230             tsock.close()
 231             name = machine_name(tivo_beacon).groups()[0]
 232         except:
 233             name = address
 234
 235         return name