beacon.py

   1 import logging
   2 import re
   3 import struct
   4 import time
   5 from socket import *
   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         if name in self.names:
  24             self.names.remove(name)
  25
  26     def addService(self, server, type, name):
  27         self.names.append(name)
  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         address = inet_aton(config.get_ip())
  37         port = int(config.getPort())
  38         for section, settings in config.getShares():
  39             ct = GetPlugin(settings['type']).CONTENT_TYPE
  40             if ct.startswith('x-container/'):
  41                 if 'video' in ct:
  42                     platform = PLATFORM_VIDEO
  43                 else:
  44                     platform = PLATFORM_MAIN
  45                 logger.info('Registering: %s' % section)
  46                 self.share_names.append(section)
  47                 desc = {'path': SHARE_TEMPLATE % quote(section),
  48                         'platform': platform, 'protocol': 'http'}
  49                 tt = ct.split('/')[1]
  50                 info = Zeroconf.ServiceInfo('_%s._tcp.local.' % tt,
  51                     '%s._%s._tcp.local.' % (section, tt),
  52                     address, port, 0, 0, desc)
  53                 self.rz.registerService(info)
  54                 self.share_info.append(info)
  55
  56     def scan(self):
  57         """ Look for TiVos using Zeroconf. """
  58         VIDS = '_tivo-videos._tcp.local.'
  59         names = []
  60
  61         # Get the names of servers offering TiVo videos
  62         browser = Zeroconf.ServiceBrowser(self.rz, VIDS, ZCListener(names))
  63
  64         # Give them half a second to respond
  65         time.sleep(0.5)
  66
  67         # Now get the addresses -- this is the slow part
  68         for name in names:
  69             info = self.rz.getServiceInfo(VIDS, name)
  70             if info and 'TSN' in info.properties:
  71                 tsn = info.properties['TSN']
  72                 address = inet_ntoa(info.getAddress())
  73                 config.tivos[tsn] = address
  74                 name = name.replace('.' + VIDS, '')
  75                 self.logger.info(name)
  76                 config.tivo_names[tsn] = name
  77
  78     def shutdown(self):
  79         self.logger.info('Unregistering: %s' % ' '.join(self.share_names))
  80         for info in self.share_info:
  81             self.rz.unregisterService(info)
  82         self.rz.close()
  83
  84 class Beacon:
  85     def __init__(self):
  86         self.UDPSock = socket(AF_INET, SOCK_DGRAM)
  87         self.UDPSock.setsockopt(SOL_SOCKET, SO_BROADCAST, 1)
  88         self.services = []
  89
  90         self.platform = PLATFORM_VIDEO
  91         for section, settings in config.getShares():
  92             ct = GetPlugin(settings['type']).CONTENT_TYPE
  93             if ct.startswith('x-container/'):
  94                 if 'music' in ct or 'photo' in ct:
  95                     self.platform = PLATFORM_MAIN
  96                     break
  97
  98         if config.get_zc():
  99             logger = logging.getLogger('pyTivo.beacon')
 100             try:
 101                 logger.info('Announcing shares...')
 102                 self.bd = ZCBroadcast(logger)
 103             except:
 104                 logger.error('Zeroconf failure')
 105                 self.bd = None
 106             else:
 107                 logger.info('Scanning for TiVos...')
 108                 self.bd.scan()
 109         else:
 110             self.bd = None
 111
 112     def add_service(self, service):
 113         self.services.append(service)
 114         self.send_beacon()
 115
 116     def format_services(self):
 117         return ';'.join(self.services)
 118
 119     def format_beacon(self, conntype, services=True):
 120         beacon = ['tivoconnect=1',
 121                   'method=%s' % conntype,
 122                   'identity=%s' % config.getGUID(),
 123                   'machine=%s' % gethostname(),
 124                   'platform=%s' % self.platform]
 125
 126         if services:
 127             beacon.append('services=' + self.format_services())
 128         else:
 129             beacon.append('services=TiVoMediaServer:0/http')
 130
 131         return '\n'.join(beacon)
 132
 133     def send_beacon(self):
 134         beacon_ips = config.getBeaconAddresses()
 135         for beacon_ip in beacon_ips.split():
 136             if beacon_ip != 'listen':
 137                 try:
 138                     self.UDPSock.sendto(self.format_beacon('broadcast'),
 139                                         (beacon_ip, 2190))
 140                 except error, e:
 141                     print e
 142
 143     def start(self):
 144         self.send_beacon()
 145         self.timer = Timer(60, self.start)
 146         self.timer.start()
 147
 148     def stop(self):
 149         self.timer.cancel()
 150         if self.bd:
 151             self.bd.shutdown()
 152
 153     def listen(self):
 154         """ For the direct-connect, TCP-style beacon """
 155         import thread
 156
 157         def server():
 158             TCPSock = socket(AF_INET, SOCK_STREAM)
 159             TCPSock.bind(('', 2190))
 160             TCPSock.listen(5)
 161
 162             while True:
 163                 # Wait for a connection
 164                 client, address = TCPSock.accept()
 165
 166                 # Accept the client's beacon
 167                 client_length = struct.unpack('!I', client.recv(4))[0]
 168                 client_message = client.recv(client_length)
 169
 170                 # Send ours
 171                 message = self.format_beacon('connected')
 172                 client.send(struct.pack('!I', len(message)))
 173                 client.send(message)
 174                 client.close()
 175
 176         thread.start_new_thread(server, ())
 177
 178     def get_name(self, address):
 179         """ Exchange beacons, and extract the machine name. """
 180         our_beacon = self.format_beacon('connected', False)
 181         machine_name = re.compile('machine=(.*)\n').search
 182
 183         try:
 184             tsock = socket()
 185             tsock.connect((address, 2190))
 186
 187             tsock.send(struct.pack('!I', len(our_beacon)))
 188             tsock.send(our_beacon)
 189
 190             length = struct.unpack('!I', tsock.recv(4))[0]
 191             tivo_beacon = tsock.recv(length)
 192
 193             tsock.close()
 194
 195             name = machine_name(tivo_beacon).groups()[0]
 196         except:
 197             name = address
 198
 199         return name