beacon.py

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