#!/usr/bin/env python3
'''
=head1 NAME

tor_

=head1 DESCRIPTION

Wildcard plugin that gathers some metrics from the Tor daemon
(https://github.com/daftaupe/munin-tor).

Derived from https://github.com/mweinelt/munin-tor

This plugin requires the stem library (https://stem.torproject.org/).

This plugin requires the GeoIP library (https://www.maxmind.com) for the countries plugin.

Available plugins:

=over 4

=item tor_bandwidth   - graph the glabal bandwidth

=item tor_connections - graph the number of connexions

=item tor_countries   - graph the countries represented our connexions

=item tor_dormant     - graph if tor is dormant or not

=item tor_flags       - graph the different flags of the relay

=item tor_routers     - graph the number of routers seen by the relay

=item tor_traffic     - graph the read/written traffic

=back

=head2 CONFIGURATION

The default configuration is:

 [tor_*]
 user toranon # or any other user/group that is running tor
 group toranon
 env.torcachefile munin_tor_country_stats.json
 env.torconnectmethod port
 env.torgeoippath /usr/share/GeoIP/GeoIP.dat
 env.tormaxcountries 15
 env.torport 9051
 env.torsocket /var/run/tor/control

To make it connect through a socket, you simply need to change C<torconnectmethod>:

 env.torconnectmethod socket

=head1 COPYRIGHT

MIT License

SPDX-License-Identifier: MIT

=head1 AUTHOR

Pierre-Alain TORET <pierre-alain.toret@protonmail.com>

=head1 MAGIC MARKERS

 #%# family=auto
 #%# capabilities=autoconf suggest

=cut
'''

import collections
import json
import os
import sys

try:
    import GeoIP
    import stem
    import stem.control
    import stem.connection
    missing_dependency_error = None
except ImportError as exc:
    # missing dependencies are reported via "autoconf"
    # thus failure is acceptable here
    missing_dependency_error = str(exc)

default_torcachefile = 'munin_tor_country_stats.json'
default_torconnectmethod = 'port'
default_torgeoippath = '/usr/share/GeoIP/GeoIP.dat'
default_tormaxcountries = 15
default_torport = 9051
default_torsocket = '/var/run/tor/control'


class ConnectionError(Exception):
    """Error connecting to the controller"""


class AuthError(Exception):
    """Error authenticating to the controller"""


def authenticate(controller):
    try:
        controller.authenticate()
        return
    except stem.connection.MissingPassword:
        pass

    try:
        password = os.environ['torpassword']
    except KeyError:
        raise AuthError("Please configure the 'torpassword' "
                        "environment variable")

    try:
        controller.authenticate(password=password)
    except stem.connection.PasswordAuthFailed:
        print("Authentication failed (incorrect password)", file=sys.stderr)


def gen_controller():
    connect_method = os.environ.get('torconnectmethod', default_torconnectmethod)
    if connect_method == 'port':
        return stem.control.Controller.from_port(port=int(os.environ.get('torport',
                                                                         default_torport)))
    elif connect_method == 'socket':
        return stem.control.Controller.from_socket_file(path=os.environ.get('torsocket',
                                                                            default_torsocket))
    else:
        print("env.torconnectmethod contains an invalid value. "
              "Please specify either 'port' or 'socket'.", file=sys.stderr)
        sys.exit(1)


#########################
# Base Class
#########################


class TorPlugin(object):
    def __init__(self):
        raise NotImplementedError

    def conf(self):
        raise NotImplementedError

    @staticmethod
    def conf_from_dict(graph, labels):
        # header
        for key, val in graph.items():
            print('graph_{} {}'.format(key, val))
        # values
        for label, attributes in labels.items():
            for key, val in attributes.items():
                print('{}.{} {}'.format(label, key, val))

    @staticmethod
    def get_autoconf_status():
        try:
            import stem
        except ImportError as e:
            return 'no (failed to import the required python module "stem": {})'.format(e)
        try:
            import GeoIP  # noqa: F401
        except ImportError as e:
            return 'no (failed to import the required python module "GeoIP": {})'.format(e)
        try:
            with gen_controller() as controller:
                try:
                    authenticate(controller)
                    return 'yes'
                except stem.connection.AuthenticationFailure as e:
                    return 'no (Authentication failed: {})'.format(e)
        except stem.SocketError:
            return 'no (Connection failed)'

    @staticmethod
    def suggest():
        options = ['bandwidth', 'connections', 'countries', 'dormant', 'flags', 'routers',
                   'traffic']

        for option in options:
            print(option)

    def fetch(self):
        raise NotImplementedError


##########################
# Child Classes
##########################


class TorBandwidth(TorPlugin):
    def __init__(self):
        pass

    def conf(self):
        graph = {'title': 'Tor observed bandwidth',
                 'args': '-l 0 --base 1000',
                 'vlabel': 'bytes/s',
                 'category': 'network',
                 'info': 'estimated capacity based on usage in bytes/s'}
        labels = {'bandwidth': {'label': 'bandwidth', 'min': 0, 'type': 'GAUGE'}}

        TorPlugin.conf_from_dict(graph, labels)

    def fetch(self):
        with gen_controller() as controller:
            try:
                authenticate(controller)
            except stem.connection.AuthenticationFailure as e:
                print('Authentication failed ({})'.format(e))
                return

            # Get fingerprint of our own relay to look up the descriptor for.
            # In Stem 1.3.0 and later, get_server_descriptor() will fetch the
            # relay's own descriptor if no argument is provided, so this will
            # no longer be needed.
            fingerprint = controller.get_info('fingerprint', None)
            if fingerprint is None:
                print("Error while reading fingerprint from Tor daemon", file=sys.stderr)
                sys.exit(1)

            response = controller.get_server_descriptor(fingerprint, None)
            if response is None:
                print("Error while getting server descriptor from Tor daemon", file=sys.stderr)
                sys.exit(1)
            print('bandwidth.value {}'.format(response.observed_bandwidth))


class TorConnections(TorPlugin):
    def __init__(self):
        pass

    def conf(self):
        graph = {'title': 'Tor connections',
                 'args': '-l 0 --base 1000',
                 'vlabel': 'connections',
                 'category': 'network',
                 'info': 'OR connections by state'}
        labels = {'new': {'label': 'new', 'min': 0, 'max': 25000, 'type': 'GAUGE'},
                  'launched': {'label': 'launched', 'min': 0, 'max': 25000, 'type': 'GAUGE'},
                  'connected': {'label': 'connected', 'min': 0, 'max': 25000, 'type': 'GAUGE'},
                  'failed': {'label': 'failed', 'min': 0, 'max': 25000, 'type': 'GAUGE'},
                  'closed': {'label': 'closed', 'min': 0, 'max': 25000, 'type': 'GAUGE'}}

        TorPlugin.conf_from_dict(graph, labels)

    def fetch(self):
        with gen_controller() as controller:
            try:
                authenticate(controller)

                response = controller.get_info('orconn-status', None)
                if response is None:
                    print("No response from Tor daemon in TorConnection.fetch()", file=sys.stderr)
                    sys.exit(1)
                else:
                    connections = response.split('\n')
                    states = dict((state, 0) for state in stem.ORStatus)
                    for connection in connections:
                        states[connection.rsplit(None, 1)[-1]] += 1
                    for state, count in states.items():
                        print('{}.value {}'.format(state.lower(), count))
            except stem.connection.AuthenticationFailure as e:
                print('Authentication failed ({})'.format(e))


class TorCountries(TorPlugin):
    def __init__(self):
        # Configure plugin
        self.cache_dir_name = os.environ.get('torcachedir', None)
        if self.cache_dir_name is not None:
            self.cache_dir_name = os.path.join(
                self.cache_dir_name, os.environ.get('torcachefile', default_torcachefile))

        max_countries = os.environ.get('tormaxcountries', default_tormaxcountries)
        self.max_countries = int(max_countries)

        geoip_path = os.environ.get('torgeoippath', default_torgeoippath)
        self.geodb = GeoIP.open(geoip_path, GeoIP.GEOIP_MEMORY_CACHE)

    def conf(self):
        """Configure plugin"""

        graph = {'title': 'Tor countries',
                 'args': '-l 0 --base 1000',
                 'vlabel': 'countries',
                 'category': 'network',
                 'info': 'OR connections by state'}
        labels = {}

        countries_num = self.top_countries()

        for c, v in countries_num:
            labels[c] = {'label': c, 'min': 0, 'max': 25000, 'type': 'GAUGE'}

        TorPlugin.conf_from_dict(graph, labels)

        # If needed, create cache file at config time
        if self.cache_dir_name:
            with open(self.cache_dir_name, 'w') as f:
                json.dump(countries_num, f)

    def fetch(self):
        """Generate metrics"""
        # Fallback if cache_dir_name is not set, unreadable or any other error
        countries_num = self.top_countries()
        # If possible, read cached data instead of doing the processing twice
        if self.cache_dir_name:
            try:
                with open(self.cache_dir_name) as f:
                    countries_num = json.load(f)
            except (IOError, ValueError):
                # use the fallback value above
                pass

        for c, v in countries_num:
            print("%s.value %d" % (c, v))

    @staticmethod
    def _gen_ipaddrs_from_statuses(controller):
        """Generate a sequence of ipaddrs for every network status"""
        for desc in controller.get_network_statuses():
            ipaddr = desc.address
            yield ipaddr

    @staticmethod
    def simplify(cn):
        """Simplify country name"""
        cn = cn.replace(' ', '_')
        cn = cn.replace("'", '_')
        cn = cn.split(',', 1)[0]
        return cn

    def _gen_countries(self, controller):
        """Generate a sequence of countries for every built circuit"""
        for ipaddr in self._gen_ipaddrs_from_statuses(controller):
            country = self.geodb.country_name_by_addr(ipaddr)
            if country is None:
                yield 'Unknown'
                continue

            yield self.simplify(country)

    def top_countries(self):
        """Build a list of top countries by number of circuits"""
        with gen_controller() as controller:
            try:
                authenticate(controller)
                c = collections.Counter(self._gen_countries(controller))
                return sorted(c.most_common(self.max_countries))
            except stem.connection.AuthenticationFailure as e:
                print('Authentication failed ({})'.format(e))
                return []


class TorDormant(TorPlugin):
    def __init__(self):
        pass

    def conf(self):
        graph = {'title': 'Tor dormant',
                 'args': '-l 0 --base 1000',
                 'vlabel': 'dormant',
                 'category': 'network',
                 'info': 'Is Tor not building circuits because it is idle?'}
        labels = {'dormant': {'label': 'dormant', 'min': 0, 'max': 1, 'type': 'GAUGE'}}

        TorPlugin.conf_from_dict(graph, labels)

    def fetch(self):
        with gen_controller() as controller:
            try:
                authenticate(controller)

                response = controller.get_info('dormant', None)
                if response is None:
                    print("Error while reading dormant state from Tor daemon", file=sys.stderr)
                    sys.exit(1)
                print('dormant.value {}'.format(response))
            except stem.connection.AuthenticationFailure as e:
                print('Authentication failed ({})'.format(e))


class TorFlags(TorPlugin):
    def __init__(self):
        pass

    def conf(self):
        graph = {'title': 'Tor relay flags',
                 'args': '-l 0 --base 1000',
                 'vlabel': 'flags',
                 'category': 'network',
                 'info': 'Flags active for relay'}
        labels = {flag: {'label': flag, 'min': 0, 'max': 1, 'type': 'GAUGE'} for flag in stem.Flag}

        TorPlugin.conf_from_dict(graph, labels)

    def fetch(self):
        with gen_controller() as controller:
            try:
                authenticate(controller)
            except stem.connection.AuthenticationFailure as e:
                print('Authentication failed ({})'.format(e))
                return

            # Get fingerprint of our own relay to look up the status entry for.
            # In Stem 1.3.0 and later, get_network_status() will fetch the
            # relay's own status entry if no argument is provided, so this will
            # no longer be needed.
            fingerprint = controller.get_info('fingerprint', None)
            if fingerprint is None:
                print("Error while reading fingerprint from Tor daemon", file=sys.stderr)
                sys.exit(1)

            response = controller.get_network_status(fingerprint, None)
            if response is None:
                print("Error while getting server descriptor from Tor daemon", file=sys.stderr)
                sys.exit(1)
            for flag in stem.Flag:
                if flag in response.flags:
                    print('{}.value 1'.format(flag))
                else:
                    print('{}.value 0'.format(flag))


class TorRouters(TorPlugin):
    def __init__(self):
        pass

    def conf(self):
        graph = {'title': 'Tor routers',
                 'args': '-l 0',
                 'vlabel': 'routers',
                 'category': 'network',
                 'info': 'known Tor onion routers'}
        labels = {'routers': {'label': 'routers', 'min': 0, 'type': 'GAUGE'}}
        TorPlugin.conf_from_dict(graph, labels)

    def fetch(self):
        with gen_controller() as controller:
            try:
                authenticate(controller)
            except stem.connection.AuthenticationFailure as e:
                print('Authentication failed ({})'.format(e))
                return
            response = controller.get_info('ns/all', None)
            if response is None:
                print("Error while reading ns/all from Tor daemon", file=sys.stderr)
                sys.exit(1)
            else:
                routers = response.split('\n')
                onr = 0
                for router in routers:
                    if router[0] == "r":
                        onr += 1

                print('routers.value {}'.format(onr))


class TorTraffic(TorPlugin):
    def __init__(self):
        pass

    def conf(self):
        graph = {'title': 'Tor traffic',
                 'args': '-l 0 --base 1024',
                 'vlabel': 'bytes/s',
                 'category': 'network',
                 'info': 'bytes read/written'}
        labels = {'read': {'label': 'read', 'min': 0, 'type': 'DERIVE'},
                  'written': {'label': 'written', 'min': 0, 'type': 'DERIVE'}}

        TorPlugin.conf_from_dict(graph, labels)

    def fetch(self):
        with gen_controller() as controller:
            try:
                authenticate(controller)
            except stem.connection.AuthenticationFailure as e:
                print('Authentication failed ({})'.format(e))
                return

            response = controller.get_info('traffic/read', None)
            if response is None:
                print("Error while reading traffic/read from Tor daemon", file=sys.stderr)
                sys.exit(1)

            print('read.value {}'.format(response))

            response = controller.get_info('traffic/written', None)
            if response is None:
                print("Error while reading traffic/write from Tor daemon", file=sys.stderr)
                sys.exit(1)
            print('written.value {}'.format(response))


##########################
# Main
##########################


def main():
    if len(sys.argv) > 1:
        param = sys.argv[1].lower()
    else:
        param = 'fetch'

    if param == 'autoconf':
        print(TorPlugin.get_autoconf_status())
        sys.exit()
    elif param == 'suggest':
        TorPlugin.suggest()
        sys.exit()
    else:
        if missing_dependency_error is not None:
            print("Failed to run tor_ due to missing dependency: {}"
                  .format(missing_dependency_error), file=sys.stderr)
            sys.exit(1)
        # detect data provider
        if __file__.endswith('_bandwidth'):
            provider = TorBandwidth()
        elif __file__.endswith('_connections'):
            provider = TorConnections()
        elif __file__.endswith('_countries'):
            provider = TorCountries()
        elif __file__.endswith('_dormant'):
            provider = TorDormant()
        elif __file__.endswith('_flags'):
            provider = TorFlags()
        elif __file__.endswith('_routers'):
            provider = TorRouters()
        elif __file__.endswith('_traffic'):
            provider = TorTraffic()
        else:
            print('Unknown plugin name, try "suggest" for a list of possible ones.',
                  file=sys.stderr)
            sys.exit(1)

        if param == 'config':
            provider.conf()
        elif param == 'fetch':
            provider.fetch()
        else:
            print('Unknown parameter "{}"'.format(param), file=sys.stderr)
            sys.exit(1)


if __name__ == '__main__':
    main()