#!/usr/local/bin/python3.12


__license__ = 'GPL v3'
__copyright__ = '2015, Kovid Goyal <kovid at kovidgoyal.net>'

import os
import shutil
import subprocess
import time
from collections import namedtuple

from calibre.constants import ismacos
from calibre.ptempfile import TemporaryDirectory
from calibre.srv.errors import HTTPForbidden
from calibre.srv.routes import Router, endpoint
from calibre.srv.tests.base import BaseTest, TestServer
from polyglot import http_client
from polyglot.binary import as_base64_bytes
from polyglot.builtins import iteritems, itervalues
from polyglot.http_cookie import CookieJar
from polyglot.urllib import HTTPBasicAuthHandler, HTTPCookieProcessor, HTTPDigestAuthHandler, HTTPError, build_opener

REALM = 'calibre-test'
is_ci = os.environ.get('CI', '').lower() == 'true'


@endpoint('/open', auth_required=False)
def noauth(ctx, data):
    return 'open'


@endpoint('/closed', auth_required=True)
def auth(ctx, data):
    return 'closed'


@endpoint('/android', auth_required=True, android_workaround=True)
def android(ctx, data):
    return 'android'


@endpoint('/android2', auth_required=True, android_workaround=True)
def android2(ctx, data):
    return 'android2'


def router(prefer_basic_auth=False, ban_for=0, ban_after=5):
    from calibre.srv.auth import AuthController
    return Router(itervalues(globals()), auth_controller=AuthController(
        {'testuser':'testpw', '!@#$%^&*()-=_+':'!@#$%^&*()-=_+'},
        ban_time_in_minutes=ban_for, ban_after=ban_after,
        prefer_basic_auth=prefer_basic_auth, realm=REALM, max_age_seconds=1))


def urlopen(server, path='/closed', un='testuser', pw='testpw', method='digest'):
    auth_handler = HTTPBasicAuthHandler() if method == 'basic' else HTTPDigestAuthHandler()
    url = f'http://localhost:{server.address[1]}{path}'
    auth_handler.add_password(realm=REALM, uri=url, user=un, passwd=pw)
    return build_opener(auth_handler).open(url)


def digest(un, pw, nonce=None, uri=None, method='GET', nc=1, qop='auth', realm=REALM, cnonce=None, algorithm='MD5', body=b'', modify=lambda x:None):
    'Create the payload for a digest based Authorization header'
    from calibre.srv.auth import DigestAuth
    templ = ('username="{un}", realm="{realm}", qop={qop}, method="{method}",'
    ' nonce="{nonce}", uri="{uri}", nc={nc}, algorithm="{algorithm}", cnonce="{cnonce}", response="{response}"')
    h = templ.format(un=un, realm=realm, qop=qop, uri=uri, method=method, nonce=nonce, nc=nc, cnonce=cnonce, algorithm=algorithm, response=None)
    da = DigestAuth(h)
    modify(da)
    pw = getattr(da, 'pw', pw)

    class Data:

        def __init__(self):
            self.method = method

        def peek(self):
            return body
    response = da.request_digest(pw, Data())
    return ('Digest ' + templ.format(
        un=un, realm=realm, qop=qop, uri=uri, method=method, nonce=nonce, nc=nc, cnonce=cnonce, algorithm=algorithm, response=response)).encode('ascii')


class TestAuth(BaseTest):

    def test_basic_auth(self):  # {{{
        'Test HTTP Basic auth'
        r = router(prefer_basic_auth=True)
        with TestServer(r.dispatch) as server:
            r.auth_controller.log = server.log
            conn = server.connect()
            conn.request('GET', '/open')
            r = conn.getresponse()
            self.ae(r.status, http_client.OK)
            self.ae(r.read(), b'open')

            conn.request('GET', '/closed')
            r = conn.getresponse()
            self.ae(r.status, http_client.UNAUTHORIZED)
            self.ae(r.getheader('WWW-Authenticate'), f'Basic realm="{REALM}"')
            self.assertFalse(r.read())
            conn.request('GET', '/closed', headers={'Authorization': b'Basic ' + as_base64_bytes(b'testuser:testpw')})
            r = conn.getresponse()
            self.ae(r.read(), b'closed')
            self.ae(r.status, http_client.OK)
            self.ae(b'closed', urlopen(server, method='basic').read())
            self.ae(b'closed', urlopen(server, un='!@#$%^&*()-=_+', pw='!@#$%^&*()-=_+', method='basic').read())

            def request(un='testuser', pw='testpw'):
                conn.request('GET', '/closed', headers={'Authorization': b'Basic ' + as_base64_bytes(f'{un}:{pw}')})
                r = conn.getresponse()
                return r.status, r.read()

            warnings = []
            server.loop.log.warn = lambda *args, **kwargs: warnings.append(' '.join(args))
            self.ae((http_client.OK, b'closed'), request())
            self.ae((http_client.UNAUTHORIZED, b''), request('x', 'y'))
            self.ae((http_client.BAD_REQUEST, b'The username or password was empty'), request('', ''))
            self.ae(1, len(warnings))
            self.ae((http_client.UNAUTHORIZED, b''), request('testuser', 'y'))
            self.ae((http_client.BAD_REQUEST, b'The username or password was empty'), request('testuser', ''))
            self.ae((http_client.BAD_REQUEST, b'The username or password was empty'), request(''))
            self.ae((http_client.UNAUTHORIZED, b''), request('asf', 'testpw'))
    # }}}

    def test_library_restrictions(self):  # {{{
        from calibre.db.legacy import create_backend
        from calibre.srv.handler import Handler
        from calibre.srv.opts import Options
        opts = Options(userdb=':memory:')
        Data = namedtuple('Data', 'username')
        with TemporaryDirectory() as base:
            l1, l2, l3 = (os.path.join(base, 'l' + x) for x in '123')
            for l in (l1, l2, l3):
                create_backend(l).close()
            ctx = Handler((l1, l2, l3), opts).router.ctx
            um = ctx.user_manager

            def get_library(username=None, library_id=None):
                ans = ctx.get_library(Data(username), library_id=library_id)
                return os.path.basename(ans.backend.library_path)

            def library_info(username=None):
                lmap, defaultlib = ctx.library_info(Data(username))
                lmap = {k:os.path.basename(v) for k, v in iteritems(lmap)}
                return lmap, defaultlib

            self.assertEqual(get_library(), 'l1')
            self.assertEqual(library_info()[0], {f'l{i}':f'l{i}' for i in range(1, 4)})
            self.assertEqual(library_info()[1], 'l1')
            self.assertRaises(HTTPForbidden, get_library, 'xxx')
            um.add_user('a', 'a')
            self.assertEqual(library_info('a')[0], {f'l{i}':f'l{i}' for i in range(1, 4)})
            um.update_user_restrictions('a', {'blocked_library_names': ['L2']})
            self.assertEqual(library_info('a')[0], {f'l{i}':f'l{i}' for i in range(1, 4) if i != 2})
            um.update_user_restrictions('a', {'allowed_library_names': ['l3']})
            self.assertEqual(library_info('a')[0], {f'l{i}':f'l{i}' for i in range(1, 4) if i == 3})
            self.assertEqual(library_info('a')[1], 'l3')
            self.assertRaises(HTTPForbidden, get_library, 'a', 'l1')
            self.assertRaises(HTTPForbidden, get_library, 'xxx')

    # }}}

    def test_digest_auth(self):  # {{{
        'Test HTTP Digest auth'
        from calibre.srv.http_request import normalize_header_name
        from calibre.srv.utils import parse_http_dict
        r = router()
        with TestServer(r.dispatch) as server:
            r.auth_controller.log = server.log

            def test(conn, path, headers={}, status=http_client.OK, body=b'', request_body=b''):
                conn.request('GET', path, request_body, headers)
                r = conn.getresponse()
                self.ae(r.status, status)
                self.ae(r.read(), body)
                return {normalize_header_name(k):v for k, v in r.getheaders()}
            conn = server.connect()
            test(conn, '/open', body=b'open')
            auth = parse_http_dict(test(conn, '/closed', status=http_client.UNAUTHORIZED)['WWW-Authenticate'].partition(' ')[2])
            nonce = auth['nonce']
            auth = parse_http_dict(test(conn, '/closed', status=http_client.UNAUTHORIZED)['WWW-Authenticate'].partition(' ')[2])
            self.assertNotEqual(nonce, auth['nonce'], 'nonce was re-used')
            self.ae(auth['realm'], REALM)
            self.ae(auth['algorithm'], 'MD5')
            self.ae(auth['qop'], 'auth')
            self.assertNotIn('stale', auth)
            args = auth.copy()
            args['un'], args['pw'], args['uri'] = 'testuser', 'testpw', '/closed'

            def ok_test(conn, dh, **args):
                args['body'] = args.get('body', b'closed')
                return test(conn, '/closed', headers={'Authorization':dh}, **args)

            ok_test(conn, digest(**args))
            # Check that server ignores repeated nc values
            ok_test(conn, digest(**args))

            warnings = []
            server.loop.log.warn = lambda *args, **kwargs: warnings.append(' '.join(args))
            # Check stale nonces
            orig, r.auth_controller.max_age_seconds = r.auth_controller.max_age_seconds, -1
            auth = parse_http_dict(test(conn, '/closed', headers={
                'Authorization':digest(**args)},status=http_client.UNAUTHORIZED)['WWW-Authenticate'].partition(' ')[2])
            self.assertIn('stale', auth)
            r.auth_controller.max_age_seconds = orig
            ok_test(conn, digest(**args))

            def fail_test(conn, modify, **kw):
                kw['body'] = kw.get('body', b'')
                kw['status'] = kw.get('status', http_client.UNAUTHORIZED)
                args['modify'] = modify
                return test(conn, '/closed', headers={'Authorization':digest(**args)}, **kw)

            # Check modified nonce fails
            fail_test(conn, lambda da: setattr(da, 'nonce', 'xyz'))
            fail_test(conn, lambda da: setattr(da, 'nonce', 'x' + da.nonce))

            # Check mismatched uri fails
            fail_test(conn, lambda da: setattr(da, 'uri', '/'))
            fail_test(conn, lambda da: setattr(da, 'uri', '/closed2'))
            fail_test(conn, lambda da: setattr(da, 'uri', '/closed/2'))

            # Check that incorrect user/password fails
            fail_test(conn, lambda da: setattr(da, 'pw', '/'))
            fail_test(conn, lambda da: setattr(da, 'username', '/'))
            fail_test(conn, lambda da: setattr(da, 'username', ''))
            fail_test(conn, lambda da: setattr(da, 'pw', ''))
            fail_test(conn, lambda da:(setattr(da, 'pw', ''), setattr(da, 'username', '')))

            # Check against python's stdlib
            self.ae(urlopen(server).read(), b'closed')

            # Check using curl
            curl = shutil.which('curl')
            if curl and not (is_ci and ismacos):  # curl mysteriously returns b'' in CI with no errors
                def docurl(data, *args):
                    cmd = [curl, '--silent'] + list(args) + [f'http://localhost:{server.address[1]}/closed']
                    p = subprocess.Popen(cmd, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                    stdout, stderr = p.communicate()
                    p.wait()
                    self.ae(stdout, data, f'stderr:\n{stderr.decode(errors="replace")}')
                docurl(b'')
                docurl(b'', '--digest', '--user', 'xxxx:testpw')
                docurl(b'', '--digest', '--user', 'testuser:xtestpw')
                docurl(b'closed', '--digest', '--user', 'testuser:testpw')
    # }}}

    def test_fail_ban(self):  # {{{
        ban_for = 0.5/60.0
        r = router(prefer_basic_auth=True, ban_for=ban_for, ban_after=2)
        with TestServer(r.dispatch) as server:
            r.auth_controller.log = server.log
            conn = server.connect()

            def request(un='testuser', pw='testpw'):
                conn.request('GET', '/closed', headers={'Authorization': b'Basic ' + as_base64_bytes(f'{un}:{pw}')})
                r = conn.getresponse()
                return r.status, r.read()

            warnings = []
            server.loop.log.warn = lambda *args, **kwargs: warnings.append(' '.join(args))
            self.ae((http_client.OK, b'closed'), request())
            self.ae((http_client.UNAUTHORIZED, b''), request('x', 'y'))
            self.ae((http_client.UNAUTHORIZED, b''), request('x', 'y'))
            self.ae(http_client.FORBIDDEN, request('x', 'y')[0])
            self.ae(http_client.FORBIDDEN, request()[0])
            time.sleep(ban_for * 60 + 0.01)
            self.ae((http_client.OK, b'closed'), request())
    # }}}

    def test_android_auth_workaround(self):  # {{{
        'Test authentication workaround for Android'
        r = router()
        with TestServer(r.dispatch) as server:
            r.auth_controller.log = server.log
            conn = server.connect()

            # First check that unauth access fails
            conn.request('GET', '/android')
            r = conn.getresponse()
            self.ae(r.status, http_client.UNAUTHORIZED)

            auth_handler = HTTPDigestAuthHandler()
            url = f'http://localhost:{server.address[1]}/android'
            auth_handler.add_password(realm=REALM, uri=url, user='testuser', passwd='testpw')
            cj = CookieJar()
            cookie_handler = HTTPCookieProcessor(cj)
            r = build_opener(auth_handler, cookie_handler).open(url)
            self.ae(r.getcode(), http_client.OK)
            cookies = tuple(cj)
            self.ae(len(cookies), 1)
            cookie = cookies[0]
            self.assertIn(':', cookie.value)
            self.ae(cookie.path, '/android')
            r = build_opener(cookie_handler).open(url)
            self.ae(r.getcode(), http_client.OK)
            self.ae(r.read(), b'android')
            # Test that a replay attack against a different URL does not work
            try:
                build_opener(cookie_handler).open(url+'2')
                assert ('Replay attack succeeded')
            except HTTPError as e:
                self.ae(e.code, http_client.UNAUTHORIZED)

    # }}}
