polyglot/aggregator.py

import pickle
import requests

from backends.schoolloop import SchoolloopBackend
from backends.haiku import HaikuBackend
from backends.webassign import WebassignBackend
from backends.dummy import DummyBackend

backends = [DummyBackend, SchoolloopBackend, HaikuBackend, WebassignBackend]

__STATE_CACHE = None

def enable_logging():
    from http.client import HTTPConnection
    import logging
    HTTPConnection.debuglevel = 1
    logging.basicConfig()
    logging.getLogger().setLevel(logging.DEBUG)
    requests_log = logging.getLogger("requests.packages.urllib3")
    requests_log.setLevel(logging.DEBUG)
    requests_log.propagate = True

def load_state(name):
    global __STATE_CACHE

    if __STATE_CACHE == None:
        try:
            with open("state.pkl", "rb") as f:
                __STATE_CACHE = pickle.load(f)
        except FileNotFoundError:
            __STATE_CACHE = {}

    try:
        return __STATE_CACHE[name]
    except KeyError:
        __STATE_CACHE[name] = {"authed": False}
        return __STATE_CACHE[name]

def save_state():
    with open("state.pkl", "wb") as f:
        pickle.dump(__STATE_CACHE, f)

def load_backend(name):
    sess = requests.Session()

    state = load_state(name)

    try:
        sess.cookies = state["cookies"]
    except KeyError:
        print("No cookies loaded")

    backend = { x.__name__ : x for x in backends } [name] ()
    backend.session = sess

    return backend

def authenticate(backend, config):
    state = load_state(backend.__class__.__name__)

    if state["authed"]:
        backend.fast_login(config)
        return True

    if backend.login(config):
        backend.authenticated = True

        state["authed"] = True
        state["cookies"] = backend.session.cookies
        save_state()
        
        return True
    else:
        print("Bad credentials")
        return False