keenetic.py - pyscript for Home Assistant

"""
The script allows to modify internet availability schedule for particular device connected to Keenetic router.

Script to be used with Pyscript HACS custom integration (https://github.com/custom-components/pyscript).

The following configuration is required (at configuration.yaml):
    pyscript:
      allow_all_imports: true
      apps:
        keenetic:
          - ip: 192.168.1.1
            login: !secret keenetic_login
            pass: !secret keenetic_pass

Binary sensor could be created to reflect connection status (changed only when connection is changed from HA):
    mqtt:
      binary_sensor:
        - state_topic: "homeassistant/binary_sensor/any_topic_name"
          name: "any_name"
          payload_off: "off"
          payload_on: "on"

The following service could be used to call the script:
    service: pyscript.keenetic_apply_schedule
    data:
      mac: '11:22:33:44:55:66' # mac address of the device to apply new schedule
      schedule: 'schedule0' # schedule id or False to disable schedule
      mqtt_topic: 'homeassistant/binary_sensor/any_topic_name'
"""

import aiohttp
import hashlib

@service
def keenetic_apply_schedule(mac=None, schedule=None, mqtt_topic=None):
    log.debug(f'Starting keenetic_apply_schedule service for {mac}...')

    log.debug('Reading keenetic configuration...')
    app_config = pyscript.config.get('apps').get('keenetic', [{}])[0]
    keenetic_ip = app_config.get('ip', '192.168.0.1')
    keenetic_login = app_config.get('login', 'admin')
    keenetic_pass = app_config.get('pass', '1234')
    log.debug(f'Finished to read configuration, device ip is {keenetic_ip}')
    
    # https://forum.keenetic.com/topic/4163-%D0%B8%D1%81%D0%BF%D0%BE%D0%BB%D1%8C%D0%B7%D0%BE%D0%B2%D0%B0%D0%BD%D0%B8%D0%B5-ifttt-%D0%B8-zapier/
    # https://forum.keenetic.com/topic/5461-%D0%BD%D0%B5-%D1%80%D0%B0%D0%B1%D0%BE%D1%82%D0%B0%D0%B5%D1%82-digest-%D0%B0%D0%B2%D1%82%D0%BE%D1%80%D0%B8%D0%B7%D0%B0%D1%86%D0%B8%D1%8F-%D0%B4%D0%BB%D1%8F-rci/
    log.debug('Starting authorization process...')
    authorized = False
    url = f'http://{keenetic_ip}/auth'
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as resp:
            resp_cookies = resp.cookies
            if resp.status == 200:
                log.debug('Already authorized')
                authorized = True
            else:
                resp_token = resp.headers.get('X-NDM-Challenge', None)
                resp_realm = resp.headers.get('X-NDM-Realm', None)
                if not (resp_cookies and resp_token and resp_realm):
                    log.error('No auth data received')
                else:
                    log.debug(f'Auth data received: resp_cookies - {resp_cookies}, resp_token - {resp_token}, resp_realm - {resp_realm}')
                    
                    log.debug('Calculating hashes...')
                    part1 = resp_token.encode('utf-8')
                    log.debug(f'Part1 = {part1}')
                    for_md5 = keenetic_login + ':' + resp_realm + ':' + keenetic_pass
                    for_md5 = for_md5.encode('utf-8')
                    part2 = hashlib.md5(for_md5)
                    part2 = part2.hexdigest().encode('utf-8')
                    log.debug(f'Part2 = {part2}')
                    auth_hash = hashlib.sha256(part1 + part2).hexdigest()
                    log.debug(f'Finished to calculate the hash, the value is {auth_hash}')
                    
                    log.debug('Authorizing...')
                    req_data = {'login': keenetic_login, 'password': auth_hash}
                    # it is important to use the cookies
                    async with session.post(url, json=req_data, cookies=resp_cookies) as resp:
                        if resp.status == 200:
                            # the session is valid for 10 minutes
                            log.debug('Authorized successfully')
                            authorized = True
                        else:
                            log.error(f'Failed to authorize, the error code is {resp.status}')
                    log.debug('Finished to authorize')
            log.debug('Finished authorization process')
            
            if authorized:
                url = f'http://{keenetic_ip}/rci/ip/hotspot/host'
                req_data = {'mac': mac, "permit": True, "policy": False, "schedule": schedule}
                async with session.post(url, json=req_data, cookies=resp_cookies) as resp:
                    if resp.status == 200:
                        log.debug('New schedule applied successfully')
                        log.debug(resp.text())
                        
                        if mqtt_topic:
                            log.debug('Sending MQTT message to update sensor state')
                            if schedule:
                                # https://hacs-pyscript.readthedocs.io/en/stable/reference.html#calling-services
                                service.call('mqtt', 'publish', topic=mqtt_topic, payload="off")
                            else:
                                service.call('mqtt', 'publish', topic=mqtt_topic, payload="on")
                    else:
                        log.error(f'Failed to apply new schedule, the error code is {resp.status}')
                
            
    log.debug('Finished to execute keenetic_apply_schedule service')