#!/usr/bin/env python3

# 3rd-party dependencies: selenium >=4.10.0, geckodriver

import os
import sys

from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.firefox.service import Service

username = os.environ['ROUTER_USERNAME']
password = os.environ['ROUTER_PASSWORD']
target_ip = sys.argv[1]

PORTS = {22, 25, 80, 443, 465, 993, 51820}

options = webdriver.FirefoxOptions()
options.add_argument('-headless')
service = Service(executable_path='/usr/local/bin/geckodriver', service_log_path=os.devnull)
browser = webdriver.Firefox(options=options, service=service)
browser.delete_all_cookies()
browser.implicitly_wait(1)
browser.get('http://192.168.0.1/login.asp')

username_text_box = browser.find_element(By.ID, 'loginUsername')
username_text_box.send_keys(username)

password_text_box = browser.find_element(By.ID, 'loginPassword')
password_text_box.send_keys(password)

submit_button = browser.find_element(By.XPATH, '/html/body/div[1]/div[4]/div[1]/form/button')
submit_button.click()

browser.get('http://192.168.0.1/advanced/forwarding.asp')
form_fields = '/html/body/div/div[4]/div[1]/form/input'
forwarded_ports = browser.find_elements(By.XPATH, form_fields)
forwarded: set[int] = set()
forwarded = {int(port.get_attribute('value').split('-')[0]) for port in forwarded_ports[1::4]}

missing = PORTS - forwarded
n_missing = len(missing)
add_row = browser.find_element(By.XPATH, '/html/body/div/div[4]/div[1]/form/button[1]')
for _ in range(n_missing):
	add_row.click()

n_fwd = len(forwarded)
for i, port in enumerate(missing):
	browser.find_element(By.NAME, f'ForwardingPortRange{i + n_fwd}').clear()
	browser.find_element(By.NAME, f'ForwardingPortRange{i + n_fwd}').send_keys(str(port))
	browser.find_element(By.NAME, f'ForwardingTargetIp{i + n_fwd}').clear()
	browser.find_element(By.NAME, f'ForwardingTargetIp{i + n_fwd}').send_keys(target_ip)
	browser.find_element(By.NAME, f'ForwardingTargetPort{i + n_fwd}').clear()
	browser.find_element(By.NAME, f'ForwardingTargetPort{i + n_fwd}').send_keys(str(port))
submit = browser.find_element(By.XPATH, '/html/body/div/div[4]/div[1]/form/button[2]')
submit.click()

logout = browser.find_element(By.CLASS_NAME, 'logout')
logout.click()

browser.quit()