beacon.py

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