#!/usr/bin/env python 
import socket
from struct import *
import datetime
import pcapy
import sys
import time
from dns import ip2name

import config

group_to_vlm={}

for g in config.multicast.values():
    for tuple in g.values():
        group_to_vlm[tuple[1]]=group_to_vlm.get(tuple[1], []) + [tuple[0]]

def run_vlc(cmd):
    try:
        s=socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect(('127.0.0.1', 4212))
        ret=s.recv(65536)
        s.send('admin\n')
        #~ print cmd
        s.send(cmd+'\n')
        s.settimeout(0.3)
    
        try:
            while True:
                ret+=s.recv(65536)
        except socket.timeout:
            pass
        s.close()
        return ret
    except socket.error:
        return ""

def stop(group):
    if group_to_vlm.get(group, None):
        cmd='\n'.join('control %s stop' % name for name in group_to_vlm[group])
        run_vlc(cmd)
    
def play(group):
    if group_to_vlm.get(group, None):
        cmd='\n'.join('control %s play' % name for name in group_to_vlm[group])
        run_vlc(cmd)
        
def playing(group):
    ret=True
    if group_to_vlm.get(group, None):
        for name in group_to_vlm[group]:
            ret=ret and 'playing' in run_vlc('show %s' % name)
    return ret


class IgmpTable(object):
    
    def __init__(self, dev, ignore_src=[], debug=0, resolve_dns=False, logpath=None):
        self.CLEAN_FREQ=10
        self.MC_REPORT=300
        self.CLEAN_LAST=0
        self.STREAM_LAST=time.time()
        self.table={}
        self.debug=debug
        self.counter={}
        self.dev=dev
        self.ignore_src=ignore_src
        self.resolve_dns=resolve_dns
        self.logpath=logpath

        self.sock = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.IPPROTO_IGMP)
        self.sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 2)
 
    def run(self):
        cap = pcapy.open_live(self.dev , 65536 , 1 , self.CLEAN_FREQ * 1000)
        cap.setfilter('igmp')
     
        #start sniffing packets
        try:
            while(True) :
                try:
                    (header, packet) = cap.next()
                    tuple=self.parse_packet(packet)
                    if tuple and tuple[0][0] not in self.ignore_src:
                        # type: Membership Query (17), Membership Report (IGMPv1: 18, IGMPv2: 22, IGMPv3: 34), Leave Group (23)
                        ((src, dst), type, group) = tuple
                        if type in [18, 22, 34]:
                            #print "%s report %s" % (src, group)
                            self.report(group, src)
                        elif type == 23:
                            #print "%s leaving %s" % (src, group)
                            self.leave(group, src)
                        #else:
                        #    print type
                    self.clean()
                    self.clean_stream()
                except socket.timeout:
                    pass
        except (KeyboardInterrupt,):
            exit(0)

    def query(self, group):
        igmph_length = 8
        igmp_type = 17
        max_resp_time = 10 #self.CLEAN_FREQ * 5
        igmp_group = ''.join(str(chr(int(i))) for i in group.split('.'))
        igmp = pack('!BBBB' , igmp_type, max_resp_time, 0, 0) + igmp_group

        def carry_around_add(a, b):
            c = a + b
            return (c & 0xffff) + (c >> 16)

        def checksum(msg):
            s = 0
            for i in range(0, len(msg), 2):
                w = ord(msg[i]) + (ord(msg[i+1]) << 8)
                s = carry_around_add(s, w)
            return ~s & 0xffff

        igmp=igmp[0:2] + pack('H', checksum(igmp)) + igmp[4:] 
        self.sock.sendto(igmp, (group, 0))
        
    def leave(self, group, src):
        tp=time.time()
        if self.table.get(group, None) and self.table[group].get(src, None) and tp - self.table[group][src] < self.MC_REPORT:
            if len([ip for ip in self.table[group].keys() if tp - self.table[group][ip] < self.MC_REPORT])<=1:
                self.CLEAN_LAST=time.time() - 5
                self.table[group][src]=0
                #self.query(group) #Inutile, mcproxy le fait
                if self.debug>0:
                    self.print_table()
                self.log()
        else:
            if not self.table.get(group, None):
                self.table[group]={}
                self.counter[group]=0
            self.query(group)

    #function to parse a packet
    def parse_packet(self, packet) :
         
        #parse ethernet header
        eth_length = 14
         
        eth_header = packet[:eth_length]
        eth = unpack('!6s6sH' , eth_header)
        eth_protocol = socket.ntohs(eth[2])
     
        #Parse IP packets, IP Protocol number = 8
        if eth_protocol == 8 :
            #Parse IP header
            #take first 20 characters for the ip header
            ip_header = packet[eth_length:20+eth_length]
             
            #now unpack them :)
            iph = unpack('!BBHHHBBH4s4s' , ip_header)
     
            version_ihl = iph[0]
            version = version_ihl >> 4
            ihl = version_ihl & 0xF
     
            iph_length = ihl * 4
     
            ttl = iph[5]
            protocol = iph[6]
            s_addr = socket.inet_ntoa(iph[8]);
            d_addr = socket.inet_ntoa(iph[9]);
     
         
            #IGMP Packets
            if protocol == 2 :
                u = iph_length + eth_length
                igmph_length = 8
                igmp_header = packet[u:u+4]
                igmp_group = socket.inet_ntoa(packet[u+4:u+8])
                igmph = unpack('!BBH' , igmp_header)
                 
                igmp_type = igmph[0]
                max_resp_time = igmph[1]
                checksum = igmph[2]
                
                return ((s_addr, d_addr), igmp_type, igmp_group)
                 
                #h_size = eth_length + iph_length + igmph_length
                #data_size = len(packet) - h_size
                ##get data from the packet
                #data = packet[h_size:]
                
    def report(self, group, ip):
        try:
            if not self.table[group].get(ip, None):
                self.counter[group]+=1
                play(group)
            self.table[group][ip]=time.time()        
        except KeyError:
            self.table[group]={ip:time.time()}
            self.counter[group]=1
            play(group)

    def clean(self):
        if time.time() - self.CLEAN_LAST < self.CLEAN_FREQ:
            return
        for (group, ips) in self.table.items():
            for (ip, last) in ips.items():
                if time.time() - last > self.MC_REPORT:
                    del(self.table[group][ip])
                    self.counter[group]-=1
            if not self.table[group]:
                 del(self.table[group])
                 del(self.counter[group])
                 stop(group)
            elif not playing(group):
                play(group)
        self.CLEAN_LAST=time.time()
        if self.debug > 0:
            self.print_table()
        self.log()
        
    def clean_stream(self):
        if time.time() - self.STREAM_LAST < self.MC_REPORT:
            return
        for group in group_to_vlm.keys():
            if not self.counter.get(group, None):
                stop(group)
        self.STREAM_LAST=time.time()
            

    def print_table(self):
        print self.summarise()

    def log(self):
        if self.logpath:
            f=open(self.logpath, 'w')
            f.write(self.summarise())
            f.close()

    def summarise(self):
        l=self.table.items()
        l.sort()
        tp=time.time()
        def resolve(ip):
            if self.resolve_dns:
                return ip2name(ip)
            else:
                return ip
        ret=""
        for (group, ips) in l:
            ret+= "%s: %s\n" % (resolve(group), ', '.join(["%s (%ss)" % (resolve(k[0]), int(tp - k[1])) for k in ips.items()]) if len(ips) < 10 else "%s clients" % len(ips))
        ret+= "-"*80 + "\n"
        return ret

     
 
if __name__ == "__main__":
  IgmpTable('eth0', ignore_src=['138.231.136.88', '0.0.0.0'], logpath='/var/log/igmp.log', resolve_dns=True).run()
