#!/usr/bin/env python3

"""
This script sends 'magic packets' to wake-on-lan enabled ethernet
adapters and motherboards, in order to switch on the called PC. Be
sure to connect the NIC with the motherboard if neccesary, and
enable the WOL function in the BIOS.

The 'magic packet' consists of 6 times 0xFF followed by 16 times
the hardware address of the NIC. This sequence can be encapsulated
in any kind of packet. This script uses UDP packets.

DEBUG:
    After setting environment variable DEBUG=1, you can insert
    a codeline 
        interactive()
    to get an interactive shell at this point.
"""

import socket
import os
import sys
import configparser
import glob
import argparse
import logging
import pprint

# get a logger
log = logging.getLogger()

### some further notes
# socket flags: socket.AI_ADDRCONFIG -> gives info about IP|^IP and lIP|^lIP
# set default output dialect for MACs to netaddr.strategy.eui48.mac_unix
#
# it is possible to set a password for wakeonlan packets
#   pling could support this if we find someone who actually could use it.

def make_interactive():
    if os.environ.get('DEBUG'):
        try:
            import IPython
        except ImportError as e:
            print('please install package python3-ipython\n', file=sys.stderr)
            raise e
        return IPython.embed
    else:
        return lambda: 0;

interactive = make_interactive()


# as this module is not included in python ...
try:
    import netaddr
except ImportError as e:
    print('please install package python3-netaddr\n', file=sys.stderr)
    raise e

try:
    import pyroute2
except Exception as e:
    log.debug('Module pyroute2 does not exist')

# set possible logging Severities, we want to use
logSev = [logging.ERROR, logging.INFO, logging.DEBUG]

# keys that will be split, if they are multiline
multi_line_config_keys = ['mac']

def config_logging(args):
    """ simple logging configuration """
    severity = logSev[args.v] if args.v < len(logSev) else len(logSev)
    logging.basicConfig(format='%(levelname)6s: %(message)s',
                        level=severity)


myconfig = {}


## IPv6 Foo
def calc_mac_via_ipv6_autogen_ip(ipv6):
    """ IPv6 mostly generates IP-Addresses by via
        https://tools.ietf.org/html/rfc4291 Appendix A.
        So if we get an IP that was derived for nodes
        from IEEE 802 48-bit MACs, we can derive the MAC
        from the IP """
    if not ipv6.version == 6:
        log.debug('IP-Address %s is not of type IPv6', ipv6)
        return None
    if not ipv6_is_generated_from_IEEE802_48bit(ipv6):
        log.debug('Could not gather mac address from ipv6 as this one is not'
                  ' autogenerated', ipv6)
        return None

    #             0123456701234567012345670123456701234567012345670123456701234567
    mask_high = 0b1111111111111111111111110000000000000000000000000000000000000000
    mask_low  = 0b0000000000000000000000000000000000000000111111111111111111111111
    mask_7    = 0b0000001000000000000000000000000000000000000000000000000000000000
    # flip 7th bit
    t = ipv6.value ^ mask_7
    # cut away bit 25-40
    mac = ((t & mask_high) >> 16) | (t & mask_low)
    mac = netaddr.EUI(mac)
    mac.dialect = netaddr.strategy.eui48.mac_unix
    log.debug('derived mac address to %s', mac)
    return mac


def test_calc_mac_via_ipv6_autogen_ip():
    assert calc_mac_via_ipv6_autogen_ip(netaddr.IPAddress('2001:0:0:0:0:ff:fe00:0'))\
        == netaddr.EUI('02:00:00:00:00:00')


def ipv6_is_generated_from_IEEE802_48bit(ipv6):
    """ checks whether the IP was derived from an MAC
        input: netaddr.ip.IPAddress
        output: bool
    """
    if not ipv6.version == 6:
        return False
    # bit 25-39 have to be 1
    # bit 40 is ignored can be 0 or 1 (EUI48 or MAC48)
    #        0123456701234567012345670123456701234567012345670123456701234567
    mask = 0b0000000000000000000000001111111111111111111111100000000000000000
    return (ipv6.value & mask) >> 25 == 0b111111111111111


from collections import namedtuple, defaultdict
getaddrInfoRetT = namedtuple('getaddrInfoRetT',
                     ['family', 'type', 'proto', 'canonname', 'sockaddr'])
sockaddr6T = namedtuple('sockaddr6T', ['address', 'port', 'flow_info', 'scope_id'])
sockaddr4T = namedtuple('sockaddr6T', ['address', 'port'])

def ipv6_get_AAAA_records(hostname):
    # with these calls we want to get a IPv6 Answer regardless
    # it is disable on this host or not
    sockconfig = {
        'port'   : 0,
        'proto'  : socket.getprotobyname('udp'),
        'family' : socket.AddressFamily.AF_INET6,
        'type'   : socket.SOCK_DGRAM,
        #'flags'  : socket.AI_CANONNAME,
    }
    L = []
    try:
        L = socket.getaddrinfo(hostname, **sockconfig)
    except socket.gaierror as e:
        log.error("DNS is not working: %r", e)

    retS = set()
    if L:
        for answer in map(getaddrInfoRetT._make, L):
            sockaddr = sockaddr6T._make(answer.sockaddr)
            retS.add(sockaddr.address)
        log.debug('Found IPv6-Addresses for %s: %r', hostname, list(retS))
    else:
        log.debug('No IPv6-Addresses found for %s.', hostname)
    return retS

def ipv6_get_MACs(host):
    log.debug('Try to derive MAC for %r from DNS', host)
    addrS = ipv6_get_AAAA_records(host)
    macS = set()
    for addr in addrS:
        ip = netaddr.IPAddress(addr)
        if ipv6_is_generated_from_IEEE802_48bit(ip):
            mac = calc_mac_via_ipv6_autogen_ip(ip)
            macS.add(mac)
    if not macS:
        log.info('No MACs found for %s in DNS', host)
        return macS
    log.debug('Found macs for %s via IPv6 DNS: %r', host, list(macS))
    return macS

def prepare_magic_packetcontents(mac):
    # Prepare the magic packet
    send_data =  b''.fromhex('ffffffffffff') + 16 * mac.packed
    return send_data

def get_link_localNetworks(version):
    """ gather a Networks/prefix from configured Addresses
    """
    if version == 4:
        _version = 2
    elif version == 6:
        _version = 10
    else:
        raise Exception('IPAddress version must be 4 or 6')
    pyr = pyroute2.IPRoute()
    a = pyr.get_addr()
    ifnetworkS = set()
    for c in a:
        if c.get('family',0) != 2:
            continue
        for key, value in c.get('attrs', []):
            if key == 'IFA_ADDRESS':
                myIP = netaddr.IPAddress(value)
        if myIP.is_loopback():
            continue
        prefixlen = c.get('prefixlen')
        if prefixlen:
            net = netaddr.IPNetwork('%s/%s'%(myIP, prefixlen))
            ifnetworkS.add(net)
    return ifnetworkS


def derive_BroadcastsfromIPs4(ipS):
    # this function is only useful for IPv4
    # if pyroute2 is not loaded - we cannot use this funktion
    if not sys.modules.get('pyroute2'):
        return set()
    matchingBroadcastSet = set()
    if not ipS:
        return matchingBroadcastSet
    # get a set of Networks we are directly connected to
    connected_networks = get_link_localNetworks(version=4)
    for ipString in ipS:
        ip = netaddr.IPAddress(ipString)
        for network in connected_networks:
            if ip in network:
                matchingBroadcastSet.add(network.broadcast)
    return matchingBroadcastSet


def send_packet_IPv4(host, packet):
    # check if we could send something via IPv4
    sentPackets = 0
    sockconfig = {
        'port'   : 7,
        'proto'  : socket.getprotobyname('udp'),
        'family' : socket.AddressFamily.AF_INET,
        'type'   : socket.SOCK_DGRAM,
        'flags'  : socket.AI_ADDRCONFIG | socket.AI_CANONNAME,
    }
    try:
        getaddrInfoRetTL = socket.getaddrinfo(host, **sockconfig)
        ipS = set()
        if getaddrInfoRetTL:
            for answer in map(getaddrInfoRetT._make, getaddrInfoRetTL):
                sockaddr = sockaddr4T._make(answer.sockaddr)
                ipS.add(netaddr.IPAddress(sockaddr.address))
                log.debug('Found Addresses for %s: %r', host, list(ipS))
    except Exception as e:
        log.debug('IPv4 UDP send magic PacketERROR: %s', e)
        return 0
    # Broadcast it to the LAN.
    sock = socket.socket(sockconfig['family'],
                         sockconfig['type'],
                         sockconfig['proto'])
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    #binding to an interface needs root-access. :-(
    #ifno = socket.if_nametoindex('ptxnet')
    #sock.setsockopt(socket.SOL_SOCKET, SO_BINDTODEVICE, ifno)
    # instead we ask netlink for the broadcast address of interfaces
    # and check which broadcastaddress to use
    broadcastS = derive_BroadcastsfromIPs4(ipS)
    if broadcastS:
        for broadcast in broadcastS:
            sentB = sock.sendto(packet, (str(broadcast), 7))
            log.debug('Send %d Bytes to %s via IPv4 %s', sentB, host, broadcast)
            sentPackets += 1
    else:
        # try fallback to global broadcast
        # this only works, if the host has only one IPv4 Interface
        # but sometimes we are lucky
        sentB = sock.sendto(packet, ('255.255.255.255', 7))
        log.debug('Fallback: Send %d Bytes to %s via IPv4 \'255.255.255.255\'', sentB, host)
        sentPackets += 1
    return sentPackets

def getIPv6CapableInterfaces():
    # this function is only useful for IPv6
    # if pyroute2 is not loaded - we cannot use this funktion
    if not sys.modules.get('pyroute2'):
        return set()
    matchingInterfacesSet = set()

    pyr = pyroute2.IPRoute()
    a = pyr.get_links(family=socket.AF_INET6)
    interfaceS = set()
    for c in a:
        if c.get('family', 0) != 10:
            continue
        for key, value in c.get('attrs', []):
            if key == 'IFLA_IFNAME':
                myInterface = value
            if key == 'IFLA_OPERSTATE':
                myState = value
        if myInterface == 'lo':
            continue
        if myState == 'DOWN':
            continue
        interfaceS.add(myInterface)
    log.debug('Found %s IPv6 capable interfaces: %s', len(interfaceS), interfaceS)
    return interfaceS

def send_packet_IPv6(host, packet, interfaces=None):
    # Broadcast it to the LAN.
    sentPackets = 0
    if interfaces is None:
        interfaces = getIPv6CapableInterfaces()
    for interface in interfaces:
        sock = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_DONTROUTE, 1)
        sentB = sock.sendto(packet, ('ff02::1%{}'.format(interface), 7))
        if sentB >= 102:
            log.debug('Sent magic packet to interface %s', interface)
            sentPackets += 1
    return sentPackets


def loadConfig(args, config=None, confD=None):
    """ Read in the Configuration file to get CDN specific settings

    """
    if type(config) is not configparser.ConfigParser:
        config = configparser.ConfigParser(strict=False)
    if confD is None:
        confD = defaultdict(dict)

    config_files_list = set()
    for s in args.config.split():
        config_files_list.update(glob.iglob(s))

    read_file = config.read(config_files_list)
    log.debug('Read configfile %r', config_files_list)
    for section_name, section in config.items():
        for key,value in section.items():
            if key in multi_line_config_keys:
                confD[section_name][key] = value.split('\n')
            else:
                confD[section_name][key] = value 

    return confD


def list_hosts(config):
    print('Configured Hosts:')
    for i in config:
        if i != 'General':
            print('\t', i)
    print('\n')


def parse_args():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument('--list', '-l', action="store_true",
                        help='List configured hosts')
    parser.add_argument('HOSTNAME', nargs='?',
                        help="Host to wakeup / send magic packets to")
    parser.add_argument('-v', action='count',
                        help='be more verbose. You can use it %s times' %
                        (len(logSev)-1), default=0)
    parser.add_argument('--config', '-c',
                        default='/etc/pling/*.conf',
                        help="use this config file(s), default is \'/etc/pling/*.conf\'")
    parser.add_argument('--dump-config', '-d', action='store_true',
                        help="dump complete config")
    args = parser.parse_args()
    return args, parser


if __name__ == '__main__':
    args, parser = parse_args()
    config_logging(args)
    log.debug('found args: %s', args)

    myconfig = loadConfig(args)
    try:
        if args.dump_config:
            pprint.pprint(myconfig)
            sys.exit(0)

        # Use macaddresses with any seperators.
        if args.list:
            log.debug('Listing configured hosts')
            list_hosts(myconfig)
            sys.exit(0)

        # try to gather mac-addresses for this host
        macS = set()
        ## from config
        # get mac as a string
        macL = myconfig[args.HOSTNAME].get('mac')
        if macL:
            for macString in macL:
                try:
                    mac = netaddr.EUI(macString)
                    mac.dialect = netaddr.strategy.eui48.mac_unix
                except Exception as e:
                    log.debug('Error to import mac-address section<%s> valueof<mac>', host)
                macS.add(netaddr.EUI(mac))
                log.debug('Found Mac for %s in config: %s', args.HOSTNAME, mac)
        if not macS:
            log.info('No sections found for %s in config', args.HOSTNAME)
        ## from dns        
        macS.update(ipv6_get_MACs(args.HOSTNAME))
        log.info('With all algorithms: found Macs for %s: %s', args.HOSTNAME, macS)

        if not macS:
            log.error('Could not gather the mac address for host %s', args.HOSTNAME)

        packets_send = 0
        interfaces = getIPv6CapableInterfaces()
        for mac in macS:
            # prepare the magic packet
            magicpacket = prepare_magic_packetcontents(mac)

            # now try to send a magic packet to these mac-addresses
            packets_send += send_packet_IPv4(args.HOSTNAME, magicpacket)
            packets_send += send_packet_IPv6(args.HOSTNAME, magicpacket, interfaces)
        log.info('Sent %s packets to wake up %s', packets_send, args.HOSTNAME)

        if packets_send < 1:
            log.error('Could not sent magic packet for waking up %s', args.HOSTNAME)
            sys.exit(1)
    except Exception as e:
        log.error(parser.format_usage())
        raise(e)
