psk_change

#!/usr/bin/env python3
# vim: set encoding=utf8
"""
Changes PSK for specific SSID
and makes HTML file containing new SSID
"""

import sys
import random
import string
import json
import urllib
import datetime
import locale
from pymongo import MongoClient
from unifi.controller import Controller
from bson.objectid import ObjectId

HOSTNAME = '127.0.0.1'
PORT = 27117
STATIC_TEMPLATE = 'password_template.html'
STATIC_FILE = '/var/www/wifi/index.html'
DB_NAME = 'ace'
DB_SSID = 'SPECIFIC_SSID'
UNIFI_PASSWORD='password'
UNIFI_LOGIN='admin'

def generate_psk(max_len=8):
    """
    Generates pseudorandom PSK
    """
    psk = ''
    for counter in range(max_len):
        psk += ''.join(random.SystemRandom().choice(string.ascii_uppercase + string.digits))

    return psk

def generate_static(psk):
    """
    Makes HTML file with PSK
    """
    with open(STATIC_TEMPLATE, 'r') as static_template_fd:
        template_string = ''.join(static_template_fd.readlines())
        template = string.Template(template_string)
        static_template_fd.close()

    substitution = {'date': get_russian_locale_abbrev(),
            'psk': psk}
    result = template.substitute(substitution)

    with open(STATIC_FILE, 'w') as static_fd:
        static_fd.write(result)
        static_fd.close()

def update_psk2(psk):
    """
    Changes PSK in MongoDB
    """
    try:
        client = MongoClient(HOSTNAME, PORT)
        database = client[DB_NAME]
        ssids = database.wlanconf
        for ssid in ssids.find({'name': DB_SSID}):
            id = ssid['_id']
            result = ssids.find_one_and_update({'_id': ObjectId(id)},
                                               {'$set':
                                                   {'x_passphrase': psk }})
        if result:
            client.close()
            return True
        print('Error during PSK update')
        sys.exit(1)
    except Exception as error:
        print('Error during PSK update', error)
        sys.exit(1)

def reprovision_aps():
    """
    Forces AP reprovision
    """
    c = Controller(HOSTNAME, UNIFI_LOGIN, UNIFI_PASSWORD, 8443, 'v5', 'default')
    
    js = json.dumps({"enable_isolated_wlan":"true",
			"enabled":True,
			"key":"connectivity",
			"site_id":"default",
			"uplink_host":"",
			"uplink_type":"gateway"})
    params = urllib.urlencode({'json': js})
    answer = c._read(c.api_url + 'set/setting/connectivity', params)

    js = json.dumps({"enable_isolated_wlan":"true",
			"enabled":False,
			"key":"connectivity",
			"site_id":"default",
			"uplink_host":"",
			"uplink_type":"gateway"})
    params = urllib.urlencode({'json': js})
    answer = c._read(c.api_url + 'set/setting/connectivity', params)

def get_psk():
    """
    Doublecheck
    """
    try:
        client = MongoClient(HOSTNAME, PORT)
        database = client[DB_NAME]
        ssids = database.wlanconf
        ssid = ssids.find_one({'name': DB_SSID})
        psk = ssid['x_passphrase']

        return psk

    except Exception as error:
        print('Error during PSK get', error)
        sys.exit(1)

def get_russian_locale_abbrev():
    lc = locale.setlocale(locale.LC_TIME)
    try:
        locale.setlocale(locale.LC_TIME, 'ru_RU.utf8')
        now = datetime.datetime.now()
        return now.strftime('%B %Y')
    finally:
        locale.setlocale(locale.LC_TIME, lc)

def main():
    """
    start here
    """
    psk = generate_psk()
    update_psk2(psk)
    reprovision_aps()
    actual_psk = get_psk()
    generate_static(actual_psk)

if __name__ == '__main__':
    main()