#!/usr/bin/env python3

import subprocess
import requests
import argparse
import json
import random

PRIVACY_FRIENDLY_COUNTRIES = ['Sweden', 'Switzerland', 'Germany', 'Finland', 'Netherlands', 'Norway']

def run_command(command):
    result = subprocess.run(command, capture_output=True, text=True)
    if result.returncode != 0:
        raise Exception(f"Failed to execute command: {' '.join(command)}")
    return result.stdout

def get_current_exit_node():
    status = json.loads(run_command(['tailscale', 'status', '--json']))
    return status.get('Peer', {}).get('Tailnet', {}).get('ExitNode', {}).get('Name')

def verify_exit_node(exit_node):
    response = requests.get('https://am.i.mullvad.net/json')
    exit_node_hostname = response.json().get('mullvad_exit_ip_hostname')
    print(f"Current exit node hostname: {exit_node_hostname}")

    exit_node_short = exit_node.split('.')[0]
    if exit_node_short == exit_node_hostname:
        print("Exit node set successfully!")
    else:
        print("Failed to set exit node!")

def set_exit_node(exit_node=None):
    if not exit_node:
        stdout = run_command(['tailscale', 'exit-node', 'suggest'])
        exit_node = next((line.split(': ')[1].strip() for line in stdout.splitlines() if 'Suggested exit node' in line), None)
    
    print(f"Setting exit node: {exit_node}")
    run_command(['tailscale', 'set', f'--exit-node={exit_node}'])
    verify_exit_node(exit_node)

def unset_exit_node():
    run_command(['tailscale', 'set', '--exit-node='])
    print("Exit node unset successfully!")

def get_random_privacy_friendly_exit_node():
    stdout = run_command(['tailscale', 'exit-node', 'list'])
    exit_nodes = [parts[1] for parts in (line.split() for line in stdout.splitlines()) 
                  if len(parts) > 3 and parts[2] in PRIVACY_FRIENDLY_COUNTRIES]
    
    if not exit_nodes:
        raise Exception("No privacy-friendly exit nodes available")
    return random.choice(exit_nodes)

def main():
    parser = argparse.ArgumentParser(description='Manage VPN exit nodes.')
    parser.add_argument('action', nargs='?', default='start', choices=['start', 'stop', 'new', 'shh'], 
                        help='Action to perform: start (default), stop, new, or shh')
    args = parser.parse_args()

    if args.action == 'start':
        current_exit_node = get_current_exit_node()
        if current_exit_node:
            print(f"Already connected to exit node: {current_exit_node}")
        else:
            set_exit_node()
    elif args.action == 'stop':
        unset_exit_node()
    elif args.action == 'new':
        set_exit_node()
    elif args.action == 'shh':
        set_exit_node(get_random_privacy_friendly_exit_node())

if __name__ == "__main__":
    main()