flywire.py

gistfile1.txt

#!/usr/bin/env python3

import sys, json, base64, subprocess
import os.path
from os import path
from pathlib import Path
import socket
import fcntl
import struct
import ipaddress

confbase = "/etc/wireguard"
listen_port = 5555
flywire_peers = {}
wg_peers = {}
wg_node = {}
node_privkey = ""
node_pubkey = ""
node_listenport = ""
wg_range = ""
conf_ip = ""

def run_proc(command):
  try:
    proc = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
  except subprocess.CalledProcessError as err:
    print(f"ERROR: {command}")
    print(f"ERROR:", err)
    return [255, None]
  else:
    if proc.returncode != 0:
      print(f"ERROR: {command}")
      print(f"ERROR:", proc.stderr.decode('utf-8'))
    return [proc.returncode, proc.stdout.decode('utf-8')]

def get_wg_ip_range(interface):
  global wg_range
  print("Getting wireguard ip range...")
  get_ip_range = run_proc(f"consul kv get autowire/{interface}/range")
  wg_range = ipaddress.ip_network(get_ip_range[1].strip())
  print(f"  {interface} ip range: {wg_range}")
  return get_ip_range[0]

def parse_consul_data(consul_data):
  global flywire_peers
  for node in consul_data:
    autowire, device, nothing, node_ip, wgkey = node['Key'].split("/")
    if node_ip not in flywire_peers: flywire_peers[node_ip] = {}
    if node['Value'] == None:
      value = ""
    else:
      value = base64.b64decode(node['Value']).decode("utf-8")
      if wgkey == 'ip':
        value = ipaddress.ip_address(value)
    flywire_peers[node_ip][wgkey] = value
  flywire_peers_total = len(flywire_peers.keys())
  print(f"consul knows about {flywire_peers_total} peers")

def check_wireguard_conf(interface):
  global listen_port
  global conf_ip
  node_default_ip = get_default_ip_address()
  if node_default_ip in flywire_peers.keys():
    flywire_node = flywire_peers[node_default_ip]
  else:
    update_consul_kv(interface, 'port', listen_port)
    update_consul_kv(interface, 'postup', '')
    flywire_node = None


  Path(f"{confbase}/{interface}").mkdir(parents=True, exist_ok=True)

  if path.exists(f"{confbase}/{interface}/ip"):
    int_ip_file = open(f"{confbase}/{interface}/ip", "r")
    conf_ip = int_ip_file.read()
  if conf_ip.strip() == "":
    if flywire_node and flywire_node['ip'] != "":
      conf_ip = str(flywire_node['ip'])
      print(f"Using {conf_ip} as {interface} ip")
    else:
      print("Finding next available IP address for {interface}...")
      conf_ip = str(get_next_ip())
      print(f"  Using {conf_ip}")
    update_consul_kv(interface, 'ip', conf_ip)
    update_consul_kv(interface, 'allowedips', f"{conf_ip}/32")
    int_ip_file = open(f"{confbase}/{interface}/ip", "w")
    int_ip_file.write(conf_ip)
    int_ip_file.close()

  if path.exists(f"{confbase}/{interface}/private"):
    int_privkey = open(f"{confbase}/{interface}/private", "r")
    conf_privkey = int_privkey.read()
  else:
    print("Generating private key...")
    conf_privkey = generate_priv_key(interface)
    if conf_privkey != None:
      int_privkey = open(f"{confbase}/{interface}/private", "w")
      int_privkey.write(conf_privkey)
      int_privkey.close()
    else:
      return 1

  wg_conf = f"[Interface]\nListenPort = {listen_port}\nPrivateKey = {conf_privkey}"

  if path.exists(f"{confbase}/{interface}.conf"):
    int_conf = open(f"{confbase}/{interface}.conf", "r")
    if int_conf.read() != wg_conf:
      print("Config file might have errors!")
      return 1
  else:
    print("Generating config file...")
    int_conf = open(f"{confbase}/{interface}.conf", "w")
    int_conf.write(wg_conf)
    int_conf.close()

def get_next_ip():
  next_ip = next(wg_range.hosts())
  for node in flywire_peers.keys():
    if flywire_peers[node]['ip'] > next_ip:
      next_ip = flywire_peers[node]['ip'] + 1

def generate_priv_key(interface):
  gen_priv_key = run_proc("wg genkey")
  if gen_priv_key[0] == 0:
    priv_key = gen_priv_key[1].strip()
    generate_pub_key(interface, priv_key)
    return priv_key 
  else:
    return None

def generate_pub_key(interface, privkey):
  gen_pub_key = run_proc("wg genkey")
  conf_pubkey = gen_pub_key[1].strip()

def update_consul_kv(interface, key, value):
  print(f"Setting consul key 'autowire/{interface}/nodes/{get_default_ip_address()}/{key}' to '{value}'")
  set_consul_kv = run_proc(f"consul kv put autowire/{interface}/nodes/{get_default_ip_address()}/{key} {value}")

def wg_status(interface):
  wg_status = run_proc(f"ip address show {interface}")
  if wg_status[0] == 1:
    return 2
  elif wg_status[0] == 0:
    try:
      get_ip_address(interface)
    except:
      return 1
    else:
      return 0
  else:
    return wg_status[0]

def get_ip_address(ifname):
  s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  return socket.inet_ntoa(fcntl.ioctl(
    s.fileno(),
    0x8915,  # SIOCGIFADDR
    struct.pack('256s', ifname[:15].encode('utf-8'))
  )[20:24])

def wg_up(interface):
  print(f"Attempting to bring up {interface}")
  wg_up_proc = run_proc(f"wg-quick up {interface}")
  if wg_up_proc[0] != 0:
    return 1

def set_wg_int_ip(interface, ip, cidr):
  set_ip = run_proc(f"ip address add {ip}/{cidr} dev {interface}")
  if set_ip[0] != 0:
    return 1
  return 0

def wg_dump(interface):
  global wg_peers
  global node_privkey
  global node_pubkey
  global node_listenport
  
  wgdump = run_proc(f"wg show {interface} dump")
  for index,row in enumerate(wgdump[1].split('\n')):
    if index == 0:
      node_privkey,node_pubkey,node_listenport,node_fwmark = row.split('\t')
    else:
      if row != "":
        peer = row.split('\t')
        peer_ip,peer_port = peer[2].split(':')
        wg_peers[peer_ip] = {
          'pubkey': peer[0],
          'port': peer_port,
          'allowedips': peer[3]
        }
  wg_peers_total = len(wg_peers.keys())
  print(f"{interface} knows about {wg_peers_total} peers")

def get_default_ip_address():
  s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  s.connect(("8.8.8.8", 80))
  return s.getsockname()[0]

def add_peers(interface, node_ips):
  for ip in node_ips:
    node = flywire_peers[ip]
    add_wg_peer(interface, node['pubkey'], ip, node['port'], node['allowedips'])
    
def add_wg_peer(interface, pubkey, peer_public_ip, peer_port, allowed_ips):
  print(f"Adding peer {peer_public_ip}:{peer_port} {pubkey} {allowed_ips}")
  add_peer = run_proc(f"wg set {interface} peer {pubkey} endpoint {peer_public_ip}:{peer_port} allowed-ips {allowed_ips}")
  if add_peer[0] != 0:
    return 1
  return 0


interface = "wg0"
# Get the wireguard range we care about from consul
if get_wg_ip_range(interface) != 0:
  sys.exit(1)

# Load peers that are being passed in from consul
consul_data = json.load(sys.stdin)
parse_consul_data(consul_data)
check_wireguard_conf(interface)
int_status = wg_status(interface)

if int_status > 1:
  wg_up(interface)
if int_status > 0:
  set_wg_int_ip(interface, conf_ip, wg_range.prefixlen)
  if wg_status(interface) != 0:
    print(f"{interface} not configured correctly")
    sys.exit(1)
if int_status == 0:
  print(f"{interface} is configured")

# Load known peers from wireguard interface
wg_dump(interface)

new_peers = set(flywire_peers.keys()) - set(wg_peers.keys()) - set([get_default_ip_address()])
print(f"{len(new_peers)} were added to {interface}")
add_peers(interface, new_peers)