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