Zauth/src/zolfa/zauth/utils/webclient.py

from urllib.parse import urljoin

import requests
from lxml.etree import HTML

from .exceptions import *


class WebClientForm:
    def __init__(self, form):
        self.action = form.attrib['action']
        self.method = form.attrib.get('method', "GET").upper()
        self.data = {i.attrib['name']: i.attrib.get('value', "")
                     for i in form.xpath("//input")
                     if 'name' in i.attrib}

    def update(self, items):
        for k, v in items.items():
            self[k] = v

    def __getitem__(self, k):
        return self.data[k]

    def __setitem__(self, k, v):
        self.data.update({k: v})

    def __contains__(self, k):
        return k in self.data

    def __iter__(self):
        return self.data.items()

    def __repr__(self):
        return str((self.method, self.action, self.data))


class WebClient(requests.Session):
    def __init__(self, headers={}, **kwargs):
        self.url = ""
        self.res = None
        self.form = None
        super().__init__(**kwargs)
        self.headers.update(headers)

    def send(self, *args, **kwargs):
        r = super().send(*args, **kwargs)
        self.res = r
        self.url = r.url
        self.status_code = r.status_code
        self.form = None
        return r

    def find_forms(self, **filters):
        if not self.res:
            return []
        tree = HTML(self.res.content)
        filters_xpath = [f"[@{k}='{v}']" for k, v in filters.items()]
        filters_xpath = "".join(filters_xpath)
        return tree.xpath("//form" + filters_xpath)

    def select_form(self, **filters):
        forms = self.find_forms(**filters)
        if len(forms) != 1:
            if forms:
                raise TooManyFormsError(self.res, filters)
            else:
                raise FormNotFoundError(self.res, filters)
        self.form = WebClientForm(forms[0])

    def update_form(self, **arguments):
        if not self.form:
            raise NoFormSelectedError('update_form')

    def submit_form(self):
        if not self.form:
            raise NoFormSelectedError('submit_form')
        if self.form.method not in ["POST"]:
            raise NotImplementedError(
                f"Submit method '{self.form.method}' not supported.")
        self.request(
            self.form.method,
            urljoin(self.url, self.form.action),
            data=self.form.data
        )