from requests import post
from json import dumps
from json import loads
from os import environ

def login(username, password, url='http://192.168.1.1', path='/ws'):
    headers = {
        'Accept': '*/*',
        'Authorization': 'X-Sah-Login',
        'Content-Type': 'application/x-sah-ws-4-call+json'
    }

    data = {
        'service': 'sah.Device.Information',
        'method': 'createContext',
        'parameters': {
            'applicationName': 'webui',
            'username': username,
            'password': password
        }
    }

    response = post(url + path, headers=headers, data=dumps(data))

    session_key = list(dict(response.cookies).keys())[0].split('/')[0]
    session_value = list(dict(response.cookies).values())[0]
    context_id = loads(response.text)['data']['contextID']

    return session_key, session_value, context_id

def cookie(session_key, session_value, context_id):
    return f'UILang=en; lastKnownIpv6TabState=visible; {session_key}/accept-language=en-US,en; {session_key}/sessid={session_value}; sah/contextId={context_id}'

def traceroute(username, password, url='http://192.168.1.1', path='/ws'):
    session_key, session_value, context_id = login(username, password, url, path)

    headers = {
        'Accept': '*/*',
        'Content-Type': 'application/x-sah-ws-4-call+json',
        'Authorization': 'X-Sah ' + context_id,
        'Cookie': cookie(session_key, session_value, context_id)
    }

    data = {
        'service': 'Traceroute',
        'method': 'start_diagnostic',
        'parameters': {
            'host': 'www.google.com',
            'ipversion': 'IPv4'
        }
    }

    return post(url + path, headers=headers, data=dumps(data))

def reboot(username, password, url='http://192.168.1.1', path='/ws'):
    session_key, session_value, context_id = login(username, password, url, path)

    headers = {
        'Accept': '*/*',
        'Content-Type': 'application/x-sah-ws-4-call+json',
        'Authorization': 'X-Sah ' + context_id,
        'Cookie': cookie(session_key, session_value, context_id)
    }

    data = {
        'service': 'NMC',
        'method': 'reboot',
        'parameters': {
            'reason': 'WebUI reboot'
        }
    }

    return post(url + path, headers=headers, data=dumps(data))

if __name__ == '__main__':
    username = environ.get('ROUTER_USERNAME')
    password = environ.get('ROUTER_PASSWORD')
    url = environ.get('ROUTER_IP') or '192.168.1.1'
    path = environ.get('ROUTER_PATH') or '/ws'

    reboot(username, password, url, path)