#!/usr/bin/env python3
# encoding: utf-8

import subprocess
import re
import sys

# 防火墙规则的名字
# 如果你想要改名字的话, 首先不要与其他防火墙规则名字冲突
# 其次, 你可能需要手动删除用旧名字命名的防火墙规则
NETFIREWALL_RULE_NAME = "FUCK WSL2"

# WSL2虚拟网卡的名字
# 默认名为"vEthernet (WSL)".
# 如果你把WSL2虚拟网卡的名字改了, 那这里也需要改
WSL_VETHERNET_NAME = "vEthernet (WSL)"

assert sys.platform == "win32", \
    "你为什么要在 Windows 之外的系统运行这个脚本呢?"
assert subprocess.getstatusoutput("where wsl.exe 2>NUL")[0] == 0, \
    "看起来你好像没有安装 WSL..."
assert subprocess.getstatusoutput("net session 1>&2>NUL")[0] == 0, \
    "请使用管理员权限运行此脚本!"

def get_wsl_vethernet_ip():
    """
    从Windows的'ipconfig'命令返回的结果中获取WSL2虚拟网卡的ip地址
    (首选方法)
    """
    rc_, text = subprocess.getstatusoutput("ipconfig")
    assert rc_ == 0, '无法执行 "ipconfig"!'
    flag = False
    for line in text.splitlines():
        if WSL_VETHERNET_NAME in line:
            flag = True
            continue
        if line.strip() and line[0] != " ":
            flag = False
            continue
        if flag and "IPv4" in line:
            return line.strip().split(":")[1].strip()
    # 使用备选方法
    return get_wsl_vethernet_ip_backup_plan()

def get_wsl_vethernet_ip_backup_plan():
    """
    通过在wsl中执行'cat /etc/resolv.conf'命令来获取WSL2虚拟网卡的ip地址
    (备选方法, 因为需要启动wsl)
    """
    rc_, text = subprocess.getstatusoutput('wsl "cat" "/etc/resolv.conf"')
    assert rc_ == 0, '无法执行 "wsl cat /etc/resolv.conf"!'
    for line in text.splitlines():
        if "nameserver" in line:
            return line.strip().split(" ", 1)[1]
    raise Exception('无法获取 WSL2 虚拟网卡的 ip 地址!')

def is_ipv4_address(ip_):
    """ 正则校验是否为标准的ipv4地址 """
    return bool(
        re.search(r'^((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)$', ip_)
    )

def set_wsl_firewall_rule(ip_):
    """ 正式开始设置防火墙规则 """
    assert is_ipv4_address(ip_), '"%s" 不是一个标准的 ipv4 地址!' % ip_
    ip_subnet = re.sub(r'\.\d+$', ".0/24", ip_)
    # 先检查以NETFIREWALL_RULE_NAME命名的防火墙规则是否存在, 有则更新, 没有则新建
    if subprocess.getstatusoutput('powershell "Get-NetFirewallRule -DisplayName \\"%s\\"' % NETFIREWALL_RULE_NAME)[0] == 0:
        firewall_rule_func = "Set-NetFirewallRule"
    else:
        firewall_rule_func = "New-NetFirewallRule"
    return subprocess.getstatusoutput(
        'powershell '
        '"%s -DisplayName \\"%s\\" -RemoteAddress \\"%s\\" -Enabled True"'
        % (firewall_rule_func, NETFIREWALL_RULE_NAME, ip_subnet)
    )[0]

if __name__ == "__main__":
    wsl_ip = get_wsl_vethernet_ip()
    print("- WSL2 虚拟网卡 ipv4 地址: %s" % wsl_ip)
    print("- 正在更新防火墙规则...")
    rc = set_wsl_firewall_rule(wsl_ip)
    if rc == 0:
        print("- 完成!")
    sys.exit(rc)