#!/usr/bin/env python
# -*- coding: utf-8 -*-

# imports
import os, sys
import platform, ctypes
import subprocess

import argparse
import datetime, time
import re

import psutil
import ipaddress

#
# variables
#

network_lan    = '172.27.0.0/16'
network_wg_sit = '172.27.240.0/21'
network_wg_cli = '172.27.248.0/24'
network_net    = '1.0.0.1'
network_server = '172.27.210.1'

#
# non-variables
#

delay = 60

os_current  = platform.system().lower()
os_specific = {
  'common': {
    'conf_dir' : '.wireguard',
    'conf_file': 'wg0.conf'
  },
  'windows': {
    'dirs': ['ProgramData|scoop', 'HOMEPATH|scoop', 'ProgramFiles|WireGuard', 'ProgramFiles(x86)|WireGuard', 'ProgramW6432|WireGuard', 'ProgramFiles', 'ProgramFiles(x86)', 'ProgramW6432'],
    'home': ['ProgramData', 'HOMEPATH'],
    'bins': ['wireguard.exe'],
    'cmd' : {
      'start': ['/installtunnelservice'],
      'stop' : ['/uninstalltunnelservice']
    }
  },
  'linux': {
    'dirs': [],
    'home': ['HOME'],
    'bins': ['wg-quick', 'wg'],
    'cmd' : {
      'start': ['up'],
      'stop' : ['down']
    }
  }
}

ip_address_lan = ipaddress.ip_network(network_lan)
ip_address_wg = {
  'sit': ipaddress.ip_network(network_wg_sit),
  'cli': ipaddress.ip_network(network_wg_cli)
}

#
# arguments parser
#

parser = argparse.ArgumentParser(
  prog        = 'wireguard autorun',
  description = 'Small app that automatically starts or stops wireguard connection',
  epilog      = 'The configuration file needs to be saved in the home directory of a user, in a subdirectory named ".wireguard" and in the file named "wg0.conf". In case a SYSTEM user on Windows is used, HOME is not populated, so ProgramData folder is used instead.'
)
parser.add_argument('-v', '--verbose', action = 'store_true')

args = parser.parse_args()

#
# functions
#

def timestamp():
  return datetime.datetime.today().strftime('%Y-%m-%dT%H:%M:%S')

def log(msg):
  if args.verbose:
    print(f'{timestamp()} - {msg}')

def check_admin():
  # try to detect the admin account
  try:
    is_admin = os.getuid() == 0
  except AttributeError:
    is_admin = ctypes.windll.shell32.IsUserAnAdmin() != 0
  else:
    is_admin = False

  # if not admin, exit the script
  if is_admin == True:
    log('skript byl spuštěn pod účtem administrátora')
  else:
    log('skript byl spuštěn pod účtem běžného uživatele, ukončuji skript')
    sys.exit()

def check_os():
  # detect OS
  if os_current not in os_specific.keys():
    log(f'systém {os_current} není podporován, ukončuji skript')
    sys.exit()
  else:
    log(f'systém {os_current} je podporován')

def ping(host, count = '5'):
  # parameter for the number of packets
  parameter = '-n' if os_current == 'windows' else '-c'
  # building the command. Ex: "ping -c 1 google.com"
  command = ['ping', parameter, count, host]
  # ping the host
  response = subprocess.Popen(command, stdout=subprocess.PIPE).stdout.read()
  # find the % of lost packages
  try:
    # check whether host is unreachable or timed out at least half of the time
    unreachable = len(re.findall(br'unreachable', response, re.IGNORECASE))
    timed_out   = len(re.findall(br'timed out', response, re.IGNORECASE))
    if (unreachable + timed_out) >= int(count)/2:
      return False
    # check if there are any lost packets
    match = re.search(br'[0-9]+%', response).group(0)
    lost  = int(match.decode('utf-8')[:-1])
  except:
    return False
  # return true or false based on the ping output
  return lost == 0

def get_bin():
  # variables
  dirs = os_specific[os_current]['dirs']
  bins = os_specific[os_current]['bins']

  # try to find wireguard bin
  full_path = ''
  if len(dirs) == 0:
    # linux-like
    command = ['which', bins[0]]
    full_path = subprocess.Popen(command, stdout=subprocess.PIPE).stdout.read()
  else:
    # windows
    for dir in dirs:
      if '|' in dir:
        dir_first  = dir.split('|',1)[0]
        dir_second = dir.split('|',1)[1]
        if dir_first not in os.environ:
          continue
        root = os.environ[dir_first] + os.sep + dir_second
      else:
        if dir not in os.environ:
          continue
        root = os.environ[dir]

      for base, _, subfiles in os.walk(root):
        if bins[0] in subfiles:
          full_path = base + os.sep + bins[0]
          break

      if full_path != '':
        break

  # strip the output and return the path
  return full_path.strip()

def get_commands():
  commands = {}

  wireguard_bin = get_bin()
  if wireguard_bin != '':
    # continue if wireguard bin has been found
    # first find home dir and its path
    for home_env in os_specific[os_current]['home']:
      if home_env not in os.environ:
        continue
      home_path = os.environ[home_env]
      if home_path != '':
        break
    if home_path == '':
      log('domovská složka uživatele nebyla nalezena, ukončuji skript')
      sys.exit()
    else:
      log(f'domovská složka nalezena v "{home_path}"')

    # and then set paths
    conf_path = home_path + os.sep + os_specific['common']['conf_dir'] + os.sep + os_specific['common']['conf_file']
    conf_name = os_specific['common']['conf_file'].split('.')[0]

    # create a list of commands
    if os.path.exists(conf_path):
      log('konfigurační soubor wireguard existuje')
      commands['start'] = [wireguard_bin] + os_specific[os_current]['cmd']['start'] + [conf_path]
      commands['stop']  = [wireguard_bin] + os_specific[os_current]['cmd']['stop']  + [conf_name]
    else:
      log('konfigurační soubor wireguardu neexistuje, ukončuji skript')
      commands = False
      sys.exit()

  else:
    # if wireguard has not been found, exit the script
    log('wireguard není na tomto počítači instalován, ukončuji skript')
    sys.exit()

  return commands

def wireguard(task):
  # log the output
  response = subprocess.Popen(commands[task], stdout=subprocess.PIPE).stdout.read().decode('utf-8')

  # display information about any errors and so on
  if 'denied' in response:
    log('nebylo možné spustit wireguard, je nutné mít vyšší oprávnění')
    sys.exit()

def run():
  # variables
  connections   = []
  connected_net = False
  connected_lan = False

  # detect internet connection
  if ping(network_net):
    connected_net = True
    log('počítač je připojený k internetu')
  else:
    log('počítač není připojený k internetu')

  # detect LAN connection
  if ping(network_server):
    connected_lan = True
    log('počítač je připojený k vnitřní síti')
  else:
    log('počítač není připojený k vnitřní síti')

  # try to discover whether the computer is connected to the local network
  interfaces_addr = psutil.net_if_addrs()
  interfaces_stat = psutil.net_if_stats()

  for if_name, if_addresses in interfaces_addr.items():
    # skip interfaces that are not active
    if if_name not in interfaces_stat.keys() or getattr(interfaces_stat[if_name], 'isup') == False:
      continue
    # parse IP addresses
    for if_address in if_addresses:
      # find the IP addressess only
      try:
        ip_address = ipaddress.ip_address(if_address.address)
        if ip_address.version == 4 and ip_address in ip_address_lan and connected_lan:
          # if the IP address is of version 4 and if it is within the LAN range, ping the central server
          if ip_address in ip_address_wg['sit'] or ip_address in ip_address_wg['cli']:
            # if the IP is within the wireguard range, mark it that way; wireguard is connected
            log(f'wireguard je připojený do sítě na IP adrese {str(ip_address)}')
            connections.append(2)
          else:
            # there is a regular connection to lan, outside of the wireguard range
            log(f'počítač je připojený přímo k vnitřní síti s IP adresou {str(ip_address)}')
            connections.append(1)
        elif connected_lan:
          # only check for the type of address
          log(f'IP adresa {str(ip_address)} nepatří do sledovaného rozsahu IP adres')
          connections.append(0)
        else:
          # other conditions
          # log(f'počítač není připojený do vnitřní sítě, IP adresa {str(ip_address)} není analyzována')
          connections.append(0)
      except:
        # the address is not an IP address (it's likely to be a MAC address); do nothing
        # log(f'řetězec {str(if_address.address)} není IP adresou')
        pass

  # now do the magic
  if (1 in connections and 2 in connections) or not connected_net:
    # kill wireguard because the computer is connected to the local network already
    # also kill wireguard if the computer is disconnected from internet
    log('pokusím se zastavit wireguard')
    wireguard('stop')
  elif len(connections) == connections.count(0) and connected_net:
    # start and connect wireguard if computer is connected to internet but is not on the local network
    log('pokusím se spustit wireguard')
    wireguard('start')
  else:
    log('tentokrát neprovádím žádné změny')

  # wait for some time before a rerun
  log(f'čekám {str(delay)} sekund...')
  time.sleep(delay)

#
# init line
#

if __name__ == '__main__':
  check_admin()
  check_os()
  commands = get_commands()
  while True:
    run()
