#!/usr/bin/env python3

import subprocess
import os
import time
import sys
from collections import Counter

# get wifi interface name (only first one)
def get_wifi_ifname():
    # Call the iw dev command and capture its output
    try:
        result = subprocess.run(["iw dev"],
                                check=True,
                                shell=True,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE,
                                text=True)
    except subprocess.CalledProcessError as e:
        print("Error:", e.stderr)
        sys.exit(1)

    # Split the output into lines
    lines = result.stdout.split("\n")

    # Iterate over the lines and print the interface names
    for line in lines:
        #if "Interface" in line:
        if line.lstrip().startswith("Interface "):
            # Extract the interface name from the line
            interface_name = line.lstrip().split(" ")[1]
            return interface_name

    return ""

# check if already a regulatory domain is set (not 00: DFS-UNSET)
def reg_domain_needed():
    try:
        result = subprocess.run(["iw reg get"],
                                check=True,
                                shell=True,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE,
                                text=True)
    except subprocess.CalledProcessError as e:
        print("Error:", e.stderr)
        sys.exit(2)

    lines = result.stdout.split("\n")

    # Iterate over the lines
    for line in lines:
        if line.lstrip().startswith("country 00:"):
            return True

    return False

# get regulatory domains as country code list
def get_wifi_reg_domain(ifname):
    print("scan for wifi networks with interface " + ifname)
    commandstr="iw dev " + ifname + " scan"
    try:
        result = subprocess.run([commandstr],
                                check=True,
                                shell=True,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE,
                                text=True)
    except subprocess.CalledProcessError as e:
        if "resource busy" in e.stderr:
            print("Device or resource busy")
            # try again later
            return None
        if "Network is down" in e.stderr:
            print("Wifi is disabled")
            sys.exit(0)
        else:
            print("Error executing: " + commandstr + ":", e.stderr)
            sys.exit(3)

    # Split the output into lines
    lines = result.stdout.split("\n")

    cclist = []
    # Iterate over the lines and print the interface names
    for line in lines:
        #if "Interface" in line:
        if line.lstrip().startswith("Country: "):
            # Extract the interface name from the line
            cclist.append(line.lstrip().split()[1])

    return cclist

max_tries = 10
ifname = get_wifi_ifname()
while not ifname and max_tries > 0 and not os.path.isfile("/sys/class/net/" + ifname + "/address"):
    time.sleep(1)
    max_tries -= 1
    ifname = get_wifi_ifname()

# check if interface name is valid
if not os.path.isfile("/sys/class/net/" + ifname + "/address"):
    print("interface \"" + ifname + "\" name not valid")
    sys.exit(4)

if not reg_domain_needed():
    print("regulatory domain already set")
    sys.exit(0)

# listen to other access points to get country codes
max_tries = 3
cclist = get_wifi_reg_domain(ifname)
while not cclist and max_tries > 0:
    time.sleep(2)
    max_tries -= 1
    cclist = get_wifi_reg_domain(ifname)

# exit if no access point was found with country code
if not cclist:
    print("no access point with country code found")
    sys.exit(0)

# sort and count list
counterlist = Counter(cclist).most_common()

print(counterlist)

# get country code for setting regulatory domain
selected_cc = ""
if len(counterlist) == 1:
    selected_cc = counterlist[0][0]

if len(counterlist) > 1:
    print("warning: found multiple regulatory domains; " +
          "maybe someone has a badly configured access point")

    # check if most common country code is leading at least by 2
    if counterlist[0][1] >= counterlist[1][1] + 2:
        selected_cc = counterlist[0][0]
    else:
        print("can not decide which country code is correct; candidates are " +
                counterlist[0][0] + " and " + counterlist[1][0])
        sys.exit(0)

print("set country code " + selected_cc)
try:
    result = subprocess.run(["iw reg set " + selected_cc],
                            check=True,
                            shell=True,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            text=True)
except subprocess.CalledProcessError as e:
    print("Error:", e.stderr)
    sys.exit(4)

# success
sys.exit(0)
