#!/usr/bin/env python

import threading
import logging
import datetime
import time
import os
import pipes
import warnings
import optparse

import lpd
from hashutils import bintohex

# ignore deprecation warnings from scapy inclusion
warnings.filterwarnings( "ignore", category = DeprecationWarning )
# disable scapy warnings about ipv6 and shit like that
logging.getLogger("scapy.runtime").setLevel(logging.ERROR)

from scapy.all import srp,Ether,ARP,conf,sendp,ltoa,atol

class BitHammer:

    def __init__(self):

        # Seconds to delay in loops
        self.DELAY = 4

        # Verbosity. 0=errors, 1=greeting, 2=show new bans, 3=show details, 4=show each ban
        self.VERBOSITY = 2

        # List of IP's that are found from ipSearcher.  These are consumed by macResolver.
        self.ips = []
        # MACs are populated by ipResolver, poisoned by aprPoisoner
        self.macs = []

        # Network info
        self.interface    = conf.iface
        self.all_modes    = ['arp', 'block_ip', 'block_mac']
        self.mode         = "arp" # arp = classic ARP poisoning, block_ip - adds to a pf table
        self.command      = "/sbin/pfctl" # for block_ip, the table name for blocked hosts
        self.command_args = "-T add -t virusprot %s" # for block_ip, the table name for blocked hosts
        self.network      = None
        self.gateway      = None  # Gateway IP
        self.gateway_hw   = None  # Gateway MAC

        # Run forever
        self.running = True

        # Turn off scapy verbosity
        conf.verb = 0

        if self.VERBOSITY > 0:
            print('Starting BitHammer.')

        # Only as root
        if not os.geteuid() == 0:
            raise Exception( "BitHammer must be run as root" )

        # Discover network gateway

        if self.VERBOSITY > 2:
            print "Finding network gateway ..."
        # for route in conf.route.routes:
        for net, msk, gw, iface, addr in conf.route.routes:
          # found a route for given interface
          if iface == self.interface:
            network = ltoa( net )
            # compute network representation if not yet done
            if network.split('.')[0] == addr.split('.')[0]:
              bits = self.bitCount( msk )
              self.network = "%s/%d" % ( network, bits )
            # search for a valid network gateway
            if self.gateway is None and gw != '0.0.0.0':
              self.gateway = gw
   
        if self.gateway is not None and self.network is not None:
            if self.VERBOSITY > 2:
                print "Gateway is %s on network %s at interface %s." % ( self.gateway, self.network, self.interface )
        elif self.mode == "arp":
            raise Exception( "Error: network gateway not found." )

    def loadSettings(self):

        parser = optparse.OptionParser()
        parser.add_option('-v', '--verbose',        action="store", dest="verbose",      default=2,                        help='Amount of verbosity')
        parser.add_option('-i', '--iface',          action="store", dest="iface",        default=conf.iface,               help='Interface for binding')
        parser.add_option('-m', '--mode',           action="store", dest="mode",         default="arp",                    help='Mode to use:  arp=ARP Poisoning, block_ip=Block by IP using a command, block_mac=Block by IP using a command')
        parser.add_option('-c', '--command',        action="store", dest="command",      default="/sbin/pfctl",            help='Command to run to block by IP/MAC.')
        parser.add_option('-a', '--command_args',   action="store", dest="command_args", default="-T add -t virusprot %s", help='Arguments to pass to command to run to block by IP/MAC.')
        (opts, args) = parser.parse_args()

        if opts.iface:
            # Not ideal, but pre python 3.3 there is no portable way to get an interface list to verify.
            self.iface = os.path.basename(opts.iface)

        if opts.mode:
            if opts.mode not in self.all_modes:
                raise Exception( "Error: Supplied mode is invalid." )
            self.mode = opts.mode

        if opts.command:
            if not os.path.isfile(self.command):
                raise Exception( "Error: Supplied command does not exist (add arguments in -a)." )
            self.command = opts.command

        if opts.command:
            self.command_args = opts.command_args

        if opts.verbose:
            if (int(opts.verbose) < 0) or (int(opts.verbose) > 4):
                raise Exception( "Error: Supplied verbosity level is invalid." )
            self.VERBOSITY = opts.verbose

        self.ips = []
        self.macs = []

    def bitCount( self, n ):
        bits = 0
        while n:
            bits += n & 1
            n   >>= 1
        return bits

    def ipSearcher(self):

        if self.VERBOSITY > 2:
            print('- Searching for Bittorrent IPs on interface %s...') % (self.iface)

        # Continiously accept LDP broadcasts
        s = lpd.LPDSocket()
        while self.running:
            # Listen for BitTorrent LPD announcements
            data, sender = s.recv_announce()
            now = datetime.datetime.now()
            if data is not None:
                # Add them to the list
                infohash, port = data
                if self.VERBOSITY > 2:
                    print('  - Found Torrent: IP: {}: Hash: {}'.format(sender[0], bintohex(infohash)))
                self.ips.append(sender[0])
            else:
                if self.VERBOSITY > 2:
                    print('  - Invalid announce: IP: {}'.format(sender[0]))
        s.close()

    def ipResolver(self):
        if self.VERBOSITY > 2:
            print('= Resolving IPs to MACs...')
        while self.running:
            # Wait till there's an IP to search for.
            if not self.ips:
                time.sleep(self.DELAY)
            else:
                # Find all MAC's on the local network
                self.gateway_hw = None
                self.endpoints  = []
                
                # broadcast arping
                ans,unans = srp( Ether( dst = "ff:ff:ff:ff:ff:ff" ) / ARP( pdst = self.network ), 
                                 verbose = False, 
                                 filter  = "arp and arp[7] = 2", 
                                 timeout = 2, 
                                 iface_hint = self.network )

                for snd,rcv in ans:
                  if rcv.psrc == self.gateway:
                    self.gateway_hw = rcv.hwsrc
                  else:
                    self.endpoints.append( ( rcv.hwsrc, rcv.psrc ) )
                
                # Check the endpoints list for this ip
                for (macEnd, ipEnd) in self.endpoints:
                    for (ipSearch) in self.ips:
                        # If found, remove ip and add mac
                        if ipEnd == ipSearch:
                            # make list unique
                            if (macEnd, ipEnd) not in self.macs:
                                if self.VERBOSITY > 1:
                                    print('  = {} -> {}, new ban added'.format(ipEnd, macEnd))
                                self.macs.append( (macEnd, ipEnd) )
                            else:
                                if self.VERBOSITY > 2:
                                    print('  = {} -> {}, already banned'.format(ipEnd, macEnd))
                            self.ips.remove(ipSearch)
                                        
    def aprPoisoner(self):
        if self.VERBOSITY > 2:
            print('# Banning MACs in background thread.')
        while self.running:
            time.sleep(self.DELAY)
            # Poison in both directions
            #   gateway -> us -> target
            #   target  -> us -> gateway
            for (mac, ip) in self.macs:
                if self.VERBOSITY > 3:
                    print('# poison {}->{}, {}->{}'.format(ip, mac, self.gateway, self.gateway_hw))
                sendp( Ether( dst = self.gateway_hw ) / ARP( op = "who-has", psrc = ip,           pdst = self.gateway ), iface_hint = self.gateway )
                sendp( Ether( dst = mac )             / ARP( op = "who-has", psrc = self.gateway, pdst = ip ),           iface_hint = self.gateway )

    def ipBlocker(self):
        if self.VERBOSITY > 2:
            print('# Banning IP addresses in background thread using command: %s %s' % (self.command, self.command_args) )
        while self.running:
            time.sleep(self.DELAY)
            for (ip) in self.ips:
                if self.VERBOSITY > 3:
                    print('# Blocking IP Address: {}'.format(ip))
                os.system(self.command + ' ' + self.command_args % (pipes.quote(ip)))
            self.ips = []

    def macBlocker(self):
        if self.VERBOSITY > 2:
            print('# Banning MAC addresses in background thread using command: %s %s' % (self.command, self.command_args) )
        while self.running:
            time.sleep(self.DELAY)
            for (mac, ip) in self.macs:
                if self.VERBOSITY > 3:
                    print('# Blocking MAC Address: {}'.format(mac))
                os.system(self.command + ' ' + self.command_args % (pipes.quote(mac)))
            self.ips = []

    def hammer(self):

        self.threadSearch = threading.Thread(target=self.ipSearcher)
        self.threadSearch.daemon = True
        self.threadSearch.start()

        if self.mode == "arp":
            self.threadResolve = threading.Thread(target=self.ipResolver)
            self.threadResolve.daemon = True
            self.threadResolve.start()

            self.threadPoison = threading.Thread(target=self.aprPoisoner)
            self.threadPoison.daemon = True
            self.threadPoison.start()

        if self.mode == "block_ip":
            self.threadBlock = threading.Thread(target=self.ipBlocker)
            self.threadBlock.daemon = True
            self.threadBlock.start()

        if self.mode == "block_mac":
            self.threadResolve = threading.Thread(target=self.ipResolver)
            self.threadResolve.daemon = True
            self.threadResolve.start()

            self.threadBlock = threading.Thread(target=self.macBlocker)
            self.threadBlock.daemon = True
            self.threadBlock.start()

        while True:
            time.sleep(self.DELAY)


if __name__ == '__main__':

    bh = BitHammer()
    bh.loadSettings()

    bh.hammer()
