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