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