#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# This file is part of TorBox, an easy to use anonymizing router based on Raspberry Pi.
# Copyright (C) 2026 radio_24
# Contact: anonym@torbox.ch
# Website: https://www.torbox.ch
# Github: https://github.com/radio24/TorBox
#
# The code in this script is an adaption from code written by Isis <isis@patternsinthevoid.net> 0x0A6A58A14B5946ABDE18E207A3ADB67A2CDB8B35
# Co-author: Nonie689 <nonie689[at]eclipso[dot]ch>
# Copyright (C) 2022 Nonie689
# Github: https://github.com/nyxnor/onionjuggler
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
#
# DESCRIPTION
# This script generates a file named torrc.exclude-nodes with an
# ExcludeNodes lines containing all relays whose observed bandwidth is less
# than a given amount. The file will be stored in the directory, the
# script was called.
#
# SYNTAX
# sudo ./exclude-slow-tor-relays [-b <min_bandwidth_in_KB/s>] [-c <consensus_file>] [-d <data_dir> ] [ -o <torrc.exclude-slow>] [-h]

# Import the necessary modules for the script.
import argparse
import logging
import os
import sys

# Try to import the required functions from the stem library.
# If the stem library is not installed, print a message explaining how to install it.
# If some other error occurs while importing, print a message that there was an error importing stem.
try:
    from stem.descriptor import parse_file
except ImportError:
    print("""\
This script requires Stem. If you're on a Debian-based system, try installing\n
the `python-stem` package. See https://stem.torproject.org/ for more info.""")
except Exception:
    print("There was an error importing Stem.")

# The minimum bandwidth a relay should have, in kilobytes. Relays with less
# than this amount of bandwidth will be added to the ExcludeNodes line and
# therefore not used by your Tor client. (default: 4000 KB/s)
MINIMUM_BANDWIDTH = 4000

# The path to the directory which your Tor process stores data in. In your
# torrc file, this directory is given under the option `DataDirectory`. Please
# see `man (1) tor` for more information, including the default path.
TOR_DATA_DIR = "/var/lib/tor"

# The path where torrc.exclude-slow will be stored
TOR_STORE_DIR = "/etc/tor"

# The file to write the new ExcludeNodes torrc lines
FILENAME = "torrc.exclude-slow"

# The path to Tor's init.d script (default: "/etc/init.d/tor-git")

# The names of various flavours of consensus files which we might find in the
# :data:`TOR_DATA_DIR`. These files should hold basic information on all the
# relays we know about, including bandwidth information. You probably don't
# need to change this setting.
POTENTIAL_CONSENSII = ["cached-consensus", "cached-microdesc-consensus"]

# Set up logging for the script.
log = logging.getLogger()


# Parse and return command-line arguments
def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--data-dir', type=str, help=("Path to tor's data directory (default: {TOR_DATA_DIR})"))
    parser.add_argument('-c', '--consensus-file', type=str, help=("The filename of tor's consensus file (default: %s)" % POTENTIAL_CONSENSII[-1]))
    parser.add_argument('-b', '--min-bandwidth', type=int, help=("The minimum bandwidth in KB/s a relay should have (default: {MINIMUM_BANDWIDTH})"), default=MINIMUM_BANDWIDTH)
    parser.add_argument('-s', '--store-dir', type=str, help=("The path where the exclude-nodes file will be stored (default: {TOR_STORE_DIR})"), default=TOR_STORE_DIR)
    parser.add_argument('-o', '--output-file', type=str, help=("The file to write the new ExcludeNodes torrc lines (default: {FILENAME})"), default=FILENAME)
#    parser.add_argument('-i', '--in-place', help=("Edit the primary torrc file in place. Provide the path to the torrc file to edit in place"))
    args = parser.parse_args()
    return args


# Define a function to create an ExcludeNodes line for the exclude-nodes file.
def create_exclude_nodes_line(fingerprints):
    exclude_nodes = 'ExcludeNodes ' + ','.join(fingerprints) + '\n'
    return exclude_nodes


# Define a function to find slow nodes in the consensus file.
def find_slow_nodes(consensus, minimum_bandwidth):
    consensus_file = open(consensus, 'rb')
    descriptors = list(parse_file(consensus_file))
    total_relays = len(descriptors)

    consensus_file.close()

    too_slow = []

    for relay in descriptors:
        if relay.bandwidth <= minimum_bandwidth:
            too_slow.append(relay.fingerprint)
            log.debug("Excluding %s with bandwidth=%s" %
                      (relay.fingerprint, relay.bandwidth))
        elif relay.is_unmeasured:
            too_slow.append(relay.fingerprint)
            log.debug("Excluding %s with unmeasured bandwidth=%s" %
                      (relay.fingerprint, relay.bandwidth))

    too_slow = ['$%s' % fingerprint for fingerprint in too_slow]
    log.info("Excluding %s/%s relays with bandwidth <= %s KB/s."
             % (len(too_slow), total_relays, minimum_bandwidth))

    return too_slow


# Define a function to write the ExcludeNodes lines to the exclude-nodes file.
def write_torrc(data_dir, output_file, lines):
    filepath = os.path.join(os.path.abspath(data_dir), output_file)
    stat = os.stat(data_dir)
    uid = stat.st_uid
    gid = stat.st_gid

    with open(filepath, 'w') as fh:
        fh.writelines(lines)
        fh.flush()
        os.fchown(fh.fileno(), uid, gid)

    log.debug("Wrote ExcludeNodes line to new torrc file: %s" % filepath)

    return filepath


# Define a function to update the primary torrc file with the contents of the exclude-nodes file.
# This will be used with the -i --in-place option. Currently, it is disabeled for safety reasons.
def write_torrc_in_place(torrc_path, exclude):
    torrc_lines = []
    orig_exclude_line = ''

    with open(torrc_path) as fh:
        for line in fh.readlines():
            if line.lower().startswith('excludenodes'):
                orig_exclude_line = line
                excluded_nodes = line.strip().replace("ExcludeNodes ", "").split(',')
                excluded_nodes.extend(exclude)
                excluded_nodes = list(set(excluded_nodes))
                excluded_nodes.sort(reverse=True)
                log.debug("\nNew torrc Excluded Nodes: " + str(excluded_nodes))
                line = create_exclude_nodes_line(excluded_nodes)

            torrc_lines.append(line)

        if not orig_exclude_line:
            torrc_lines.append(exclude)

    with open(torrc_path, 'w') as fh:
        fh.writelines(torrc_lines)
        fh.flush()

    if orig_exclude_line:
        log.debug("Overwrote ExcludeNodes line in torrc file: %s" % torrc_path)
    else:
        log.debug("Appended ExcludeNodes line to torrc file: %s" % torrc_path)

    return torrc_path


# The main function that will be executed when the script is run.
def main():
    try:
        args = get_args()
        data_dir = args.data_dir or TOR_DATA_DIR
        store_dir = args.store_dir or TOR_STORE_DIR
        output_file = args.output_file or FILENAME
        min_bw = args.min_bandwidth or MINIMUM_BANDWIDTH
    except Exception as e:
        log.error(e)
        print("There was an error in getting the arguments. Use --help to get more information.")
        sys.exit(1)

    too_slow = []
    consensus = None
    consensii = []

    log.addHandler(logging.StreamHandler(sys.stdout))
    log.setLevel(logging.INFO)

    if args.consensus_file:
        consensii.append(args.consensus_file)
    consensii.extend(POTENTIAL_CONSENSII)

    if os.path.isdir(data_dir) and os.access(data_dir, os.R_OK):
        for filename in consensii:
            filepath = os.path.join(data_dir, filename)
            if os.path.exists(filepath) and os.path.isfile(filepath):
                if os.access(filepath, os.R_OK):
                    consensus = filepath
                    break
                else:
                    log.info("Permission denied! Can't read the consensus file in %s. Try to run it as root."
                             % (filepath))
    else:
        log.warning("Permission denied! Can't read the consensus file in %s. Try to run it as root."
                    % (data_dir))

    if not consensus:
        log.error("Could not find or read consensus file.")
        sys.exit(1)

    too_slow = find_slow_nodes(consensus, min_bw)

    if too_slow:
        exclude_nodes = create_exclude_nodes_line(too_slow)
#        if args.in_place:
#            write_torrc_in_place(args.in_place, too_slow)
        if args.output_file:
            write_torrc(store_dir, output_file, exclude_nodes)

    return exclude_nodes


# Execute the main function
if __name__ == "__main__":
    main()
