#!python

import os
import sys
import time
import json
import hmac
import hashlib
import secrets
import argparse
import dateparser

import urllib
import mimetypes
import asyncio
from aiohttp import web

import webshoes as ws
import propertybag as pb
import sparen
Log = sparen.log

import authwert


def isSafeRedirect(url, domain):
    if not url:
        return False
    # Browsers strip tab/newline/CR from URLs and normalize backslashes to
    # forward slashes before navigating. Either can smuggle a protocol-relative
    # ("//host") or absolute URL past naive validation (e.g. "/\evil.com" or
    # "\/\/evil.com" become "//evil.com"), so reject them outright.
    stripped = url.replace('\t', '').replace('\n', '').replace('\r', '')
    if '\\' in stripped:
        return False
    if stripped.startswith('//'):
        return False
    parsed = urllib.parse.urlparse(stripped)
    if not parsed.scheme and not parsed.netloc:
        return True
    return parsed.scheme in ('http', 'https') and (
        parsed.hostname == domain or
        (parsed.hostname and parsed.hostname.endswith('.' + domain))
    )


def _pwMatch(supplied, stored):
    """Constant-time password comparison that tolerates non-ASCII input.

    hmac.compare_digest raises TypeError on non-ASCII str operands, so both
    sides are encoded to UTF-8 bytes first. Any unexpected type simply fails
    the comparison rather than raising.
    """
    try:
        return hmac.compare_digest(supplied.encode('utf-8'), stored.encode('utf-8'))
    except (AttributeError, UnicodeError):
        return False


def getClientIp(req, trusted_proxies):
    """Determine the client IP to use for rate limiting.

    X-Forwarded-For is appended to by each proxy, so only the rightmost
    entries (added by proxies we control) can be trusted; anything further
    left is attacker-controlled and must never be used as a rate-limit key.

    With `trusted_proxies` reverse proxies in front of us, the real client is
    the entry `trusted_proxies` positions from the right of the chain. When
    `trusted_proxies` is 0 (the default) the X-Forwarded-For/X-Real-IP headers
    are ignored entirely and the real peer address is used instead, so a client
    connecting directly cannot spoof its address.
    """
    if trusted_proxies and trusted_proxies > 0:
        xff = req.headers.get('X-Forwarded-For', '')
        chain = [ip.strip() for ip in xff.split(',') if ip.strip()]
        if chain:
            idx = max(0, len(chain) - trusted_proxies)
            return chain[idx]
    return req.remote or 'unknown'


#-------------------------------------------------------------------
# Authentication

async def authVerify(ctx, q):

    _p = ctx.opts._p

    si = None
    try:
        if _p.cookieid in ctx.req.cookies:
            si = pb.Bag(authwert.getSessionInfo(ctx.req.cookies[_p.cookieid], _p))
    except Exception as e:
        Log(e)

    if not si:
        forwarded_uri = ctx.req.headers.get('X-Forwarded-Uri', '')
        if forwarded_uri and _p.rootpath:
            rd = urllib.parse.quote(forwarded_uri, safe='')
            return web.HTTPSeeOther(f'{_p.rootpath}/login?rd={rd}')
        r = web.Response(text='Access Denied', status=401)
        for k, v in authwert.SECURITY_HEADERS.items():
            r.headers[k] = v
        return r

    if _p.verbose:
        Log(f'Logged in: {si} : {ctx.req.path}')

    r = web.Response(text='ok', status=200)
    for k, v in authwert.SECURITY_HEADERS.items():
        r.headers[k] = v
    return r



def authResp(code, p=None, page=None, loc=None):
    if not os.path.isabs(page) if page else True:
        page = authwert.libPath(page or 'web/login.html')
    h = {}
    if loc and p:
        location = loc + '?' + urllib.parse.urlencode(p)
        Log(f'Location: {location}')
        h['Location'] = location
        h['Content-Location'] = location

    # Add security headers
    for k, v in authwert.SECURITY_HEADERS.items():
        h[k] = v

    return web.FileResponse(page, status=code, headers=h)


async def authLogo(ctx, q):
    _p = ctx.opts._p
    fpath = _p.logoimage
    if not os.path.isfile(fpath):
        return web.Response(text='Not Found', status=404)
    mime = mimetypes.guess_type(fpath)[0] or 'image/png'
    return web.FileResponse(fpath, headers={"Content-Type": mime})


async def authLogin(ctx, q):

    _p = ctx.opts._p

    if ctx.req.body_exists:
        m = await ctx.req.post()
        for k in set(m.keys()):
            q[k] = m[k]

    # Get session info
    si = None
    try:
        if _p.cookieid in ctx.req.cookies:
            si = pb.Bag(authwert.getSessionInfo(ctx.req.cookies[_p.cookieid], _p))
    except Exception as e:
        Log(e)

    # Log out user if requested
    if q.logout:
        if si:
            # Mask session ID in log (only show first 8 chars)
            sid_display = si.sid[:8] + '...' if si.sid and len(si.sid) > 8 else si.sid
            Log(f'Logging out : {si.username} : {sid_display}')
            # Remove from session store
            if si.sid and si.sid in _p.sessions:
                del _p.sessions[si.sid]
            # Add JWT to blacklist for revocation
            if hasattr(si, 'jti') and si.jti:
                if 'jwt_blacklist' not in _p:
                    _p.jwt_blacklist = {}
                _p.jwt_blacklist[si.jti] = si.exp if hasattr(si, 'exp') else (time.time() + _p.exptime)
            r = authResp(200, {'rd':q.rd}, _p.loginpage)
            r.set_cookie(_p.cookieid, '', domain=_p.domain, max_age=0, secure=_p.cookie_secure, httponly=True, samesite='Strict')
            return r

    # Is user logged in?
    if si:
        return authResp(200, {}, _p.logoutpage)

    # Default redirect
    if not q.rd:
        q.rd = _p.rootpath if _p.rootpath else '/'

    if _p.verbose:
        Log(f'Redirect: {q.rd}')

    # Is user trying to login?
    if q.username and q.password:

        # Get client IP for rate limiting. Only X-Forwarded-For entries added
        # by trusted proxies are honored; otherwise the real peer address is
        # used so the header cannot be spoofed to evade limits or to flood the
        # rate-limit maps with attacker-chosen keys.
        client_ip = getClientIp(ctx.req, _p.trusted_proxies)

        # Check rate limiting by username. This guards against credential
        # stuffing but also lets an attacker lock out a known user by name, so
        # it can be disabled with --no-rate-limit-by-username on deployments
        # where targeted lockout is the bigger concern.
        if _p.rate_limit_by_username and authwert.isRateLimited(q.username, _p):
            Log(f'Rate limited (username): {q.username}')
            r = authResp(429, {'rd':q.rd, 'error':'Too many login attempts. Please try again later.'}, None, f'{_p.rootpath}/login')
            r.headers['Retry-After'] = str(authwert.getRemainingLockoutTime(q.username, _p))
            return r

        # Check rate limiting by IP
        if authwert.isRateLimited(client_ip, _p):
            Log(f'Rate limited (IP): {client_ip}')
            r = authResp(429, {'rd':q.rd, 'error':'Too many login attempts. Please try again later.'}, None, f'{_p.rootpath}/login')
            r.headers['Retry-After'] = str(authwert.getRemainingLockoutTime(client_ip, _p))
            return r

        verified = False

        if _p.userinf and q.username in _p.userinf:
            if _pwMatch(q.password, _p.userinf[q.username]['password']):
                verified = True
        elif _p.userinf and '*' in _p.userinf:
            if _pwMatch(q.password, _p.userinf['*']['password']):
                verified = True

        # if not verified:
        if not verified and _p.dbauth:
            if _p.dbauth.verify(_p, q.username, q.password):
                verified = True

        if not verified:
            # Record failed attempt for rate limiting
            if _p.rate_limit_by_username:
                authwert.recordFailedAttempt(q.username, _p)
            authwert.recordFailedAttempt(client_ip, _p)
            Log(f'Login failed : {q.username}')
            r = authResp(403, {'rd':q.rd, 'error':'Invalid username or password'}, None, f'{_p.rootpath}/login')
            # Note: login_error cookie must NOT have httponly=True so JavaScript can read it
            r.set_cookie("login_error", "Invalid username or password", domain=_p.domain, max_age=3, secure=_p.cookie_secure, samesite='Strict')
            return r

        # Clear failed attempts on successful login
        if _p.rate_limit_by_username:
            authwert.clearFailedAttempts(q.username, _p)
        authwert.clearFailedAttempts(client_ip, _p)

        # Session id
        sid = secrets.token_urlsafe(32)

        # JWT ID for revocation
        jti = secrets.token_urlsafe(16)

        # Cookie expiration
        exp = round(time.time() + _p.exptime)

        try:

            # JWT
            if _p.prvpem:
                claims = {
                        "username"  : q.username,
                        "exp"       : exp,
                        "sid"       : sid,
                        "jti"       : jti
                    }
                cookie = authwert.createSessionToken(claims, _p)

            # Session id tokens
            else:
                cookie = sid
                _p.sessions[sid] = {'username': q.username, 'exp': exp, 'sid': sid}

            Log(f'Logging in : {_p.domain} : {q.username}')

            if not isSafeRedirect(q.rd, _p.domain):
                q.rd = _p.rootpath if _p.rootpath else '/'
            r = web.HTTPSeeOther(q.rd)
            r.set_cookie(_p.cookieid, cookie, domain=_p.domain, max_age=_p.exptime, secure=_p.cookie_secure, httponly=True, samesite='Strict')
            return r

        except Exception as e:
            Log(e)

    return authResp(200, {'rd':q.rd}, _p.loginpage)


#-------------------------------------------------------------------
# Debug
async def serveSite(ctx, q):

    _p = ctx.opts._p

    si = None
    try:
        if _p.cookieid in ctx.req.cookies:
            si = pb.Bag(authwert.getSessionInfo(ctx.req.cookies[_p.cookieid], _p))
    except Exception as e:
        Log(e)

    if not si:
        rd = urllib.parse.quote(ctx.req.path, safe='')
        return web.HTTPSeeOther(f'{_p.rootpath}/login?rd={rd}')

    p = ctx.req.path.split('/')
    while len(p) and not p[0]:
        p = p[1:]
    fname = '/'.join(p)

    p = os.path.join(_p.serve, fname)
    real_serve = os.path.realpath(_p.serve)
    real_p = os.path.realpath(p)
    if real_p != real_serve and not real_p.startswith(real_serve + os.sep):
        return web.Response(text='Forbidden', status=403)
    if not os.path.isfile(p):
        p = os.path.join(p, 'index.html')
        if os.path.isfile(p):
            return web.HTTPSeeOther(os.path.join(ctx.req.path, 'index.html'))
        return web.Response(text='Not Found', status=404)

    mime = mimetypes.guess_type(p)[0]
    if not mime:
        mime = 'text/plain'

    if _p.verbose:
        Log(f'[{mime}] {p}')

    return web.FileResponse(p, headers={"Content-Type":mime})


#-------------------------------------------------------------------
# Main

async def run(_p):

    # Web server
    _p.wsa = ws.WebShoesApp(_p.addr, _p.port, {'_p': _p, 'verbose': _p.verbose})

    # Authentication (registered first so it takes priority over the catch-all)
    _p.wsa.register('auth', 'cmd', 'q', 'evt', 'r', {
            'verify'    : authVerify,
            'login'     : authLogin,
            'logo.png'  : authLogo
        })

    # Serve catch-all: serves the directory at every path so that absolute
    # links in the site (/css, /js, etc.) resolve correctly.
    if _p.serve:
        _p.wsa.register('*', 'cmd', 'q', 'evt', 'r', {
                '*'    : serveSite
            })

    _p.wsa.start()

    # Idle
    while _p.run:
        await asyncio.sleep(3)


def main(_p):

    ap = argparse.ArgumentParser()
    ap.add_argument('--version', default=False, action='store_true', help='Show version')
    ap.add_argument('--domain', '-d', default='', type=str, help='Domain name')
    ap.add_argument('--rootpath', '-r', default='', type=str, help='Root Path')
    ap.add_argument('--buildver', '-b', default='', type=str, help='Build Version')
    ap.add_argument('--addr', '-a', default='127.0.0.1', type=str, help='Server address')
    ap.add_argument('--port', '-p', default=18401, type=int, help='Server port')
    ap.add_argument('--logdir', '-l', default='', type=str, help='Default log directory')
    ap.add_argument('--logfile', '-L', default='', type=str, help='Default log directory')
    ap.add_argument('--verbose', '-V', action='store_true', help='Verbose mode')
    ap.add_argument('--serve', '-S', default='', type=str, help='Serve a local directory with login protection')
    ap.add_argument('--cookieid', '-k', default='', type=str, help='cookieid')
    ap.add_argument('--userinf', '-u', default='', type=str, help='User information')
    ap.add_argument('--scheme', '-s', default=None, type=str, help='Network Scheme (default: http with --serve, https otherwise)')
    ap.add_argument('--authfile', default='', type=str, help='Python authorize file')
    ap.add_argument('--authparams', default='', type=str, help='Parameters to pass to auth file')
    ap.add_argument('--prvkey', default='', type=str, help='File containing private key to sign JWT tokens')
    ap.add_argument('--exptime', default=0, type=int, help='Login expire time in seconds')
    ap.add_argument('--expstr', default="after 6 days", type=str, help='Login expire time as string')
    ap.add_argument('--userlist', default=False, type=bool, help='Set to always maintain an internal list of active users')
    ap.add_argument('--algorithm', default='RS256', type=str, help='Private key algorithm')
    ap.add_argument('--loginpage', default='', type=str, help='Path to custom login page HTML file')
    ap.add_argument('--logoutpage', default='', type=str, help='Path to custom logout page HTML file')
    ap.add_argument('--logoimage', default='', type=str, help='Path to logo image file served at /auth/logo.png')
    ap.add_argument('--rate-limit-attempts', default=5, type=int, help='Max failed login attempts before lockout (default: 5)')
    ap.add_argument('--rate-limit-window', default=300, type=int, help='Rate limit window in seconds (default: 300)')
    ap.add_argument('--rate-limit-lockout', default=300, type=int, help='Lockout duration in seconds (default: 300)')
    ap.add_argument('--trusted-proxies', default=0, type=int, help='Number of trusted reverse proxies in front of authwert. When >0, the client IP for rate limiting is read from X-Forwarded-For; when 0 (default) the header is ignored and the real peer address is used (default: 0)')
    ap.add_argument('--no-rate-limit-by-username', dest='rate_limit_by_username', default=True, action='store_false', help='Disable per-username rate limiting (prevents attackers locking out a user by name; IP-based limiting still applies)')

    _p.merge(vars(ap.parse_args()))

    if _p.version:
        print(authwert.__info__["version"])
        sys.stdout.flush()  # os._exit skips buffer flushing
        os._exit(0)
        return

    if _p.logdir:
        if not os.path.isdir(_p.logdir):
            os.mkdir(_p.logdir)
        if not _p.logfile:
            _p.logfile = os.path.join(_p.logdir, 'server.log')
    if _p.logfile:
        sparen.log.setLogFile(_p.logfile)

    # Apply --serve defaults before validation so the one-liner works:
    #   authwert --userinf='...' --serve=/some/dir
    if _p.serve:
        _p.serve = os.path.realpath(os.path.expanduser(_p.serve))
        if not _p.domain:
            _p.domain = 'localhost'
        if not _p.scheme:
            _p.scheme = 'http'
        if not _p.cookieid:
            _p.cookieid = 'aw-' + hashlib.sha256(_p.serve.encode()).hexdigest()[:16]
    elif not _p.scheme:
        _p.scheme = 'https'

    # Only mark cookies Secure when actually serving over HTTPS. Forcing
    # Secure on a plain-HTTP deployment (e.g. --serve on http://localhost)
    # would prevent the browser from ever sending the cookie back.
    _p.cookie_secure = (_p.scheme == 'https')
    if not _p.cookie_secure:
        Log('WARNING: Serving over HTTP; session cookies are not marked Secure. Use HTTPS in production.')

    if not _p.domain:
        raise Exception('domain name (--domain) must be provided')

    if not _p.cookieid:
        raise Exception('Cookie id/name (--cookieid) must be provided')

    if not _p.rootpath:
        if _p.port:
            _p.rootpath = f'{_p.scheme}://{_p.domain}:{_p.port}/auth'
        else:
            _p.rootpath = f'{_p.scheme}://{_p.domain}/auth'
        # _p.rootdomain = '.'.join(_p.domain.split('.')[-2:])

    deft =  6 * 24 * 60 * 60
    maxt = 90 * 24 * 60 * 60
    if 10 > _p.exptime or maxt < _p.exptime:
        if _p.expstr:
            dt = dateparser.parse(_p.expstr)
            if dt is None:
                raise Exception(f'Could not parse --expstr value: {repr(_p.expstr)}')
            _p.exptime = int(dt.timestamp() - time.time() + 1)
        if 10 > _p.exptime or maxt < _p.exptime:
            _p.exptime = deft

    # Unique instance id
    _p.iid = secrets.token_urlsafe(32)

    # Resolve --authparams from a file if the value starts with '@'
    if _p.authparams and '@' == _p.authparams[0]:
        fpath = _p.authparams[1:]
        with open(fpath) as f:
            _p.authparams = f.read().strip()

    # Custom auth
    if _p.authfile:
        if '!' == _p.authfile[0]:
            _p.authfile = authwert.libPath(_p.authfile[1:])
        if os.path.isfile(_p.authfile):
            import importlib.util
            spec = importlib.util.spec_from_file_location("authfile",_p.authfile)
            _p.dbauth = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(_p.dbauth)
            _p.dbauth.init(_p)

    _p.loginpage = os.path.realpath(os.path.expanduser(_p.loginpage)) if _p.loginpage else authwert.libPath('web/login.html')
    _p.logoutpage = os.path.realpath(os.path.expanduser(_p.logoutpage)) if _p.logoutpage else authwert.libPath('web/logout.html')
    _p.logoimage = os.path.realpath(os.path.expanduser(_p.logoimage)) if _p.logoimage else authwert.libPath('web/authwert.png')
    if not os.path.isfile(_p.loginpage):
        raise Exception(f'Login page not found: {_p.loginpage}')
    if not os.path.isfile(_p.logoutpage):
        raise Exception(f'Logout page not found: {_p.logoutpage}')

    Log("Parameters: ", {'addr': _p.addr, 'port': _p.port, 'domain': _p.domain, 'rootpath': _p.rootpath, 'cookieid': _p.cookieid})

    _p.sessions = {}

    # Initialize rate limiting state
    _p.failed_attempts = {}
    _p.lockout_until = {}
    _p.jwt_blacklist = {}

    # Read the private key
    if _p.prvkey:
        if not authwert.readPrivateKey(_p.prvkey, _p):
            raise Exception(f'Failed to read private key from : {_p.prvkey}')

    # Get user info
    _p.userinf = _p.userinf.strip()
    if _p.userinf:
        if '{' == _p.userinf[0]:
            Log('WARNING: --userinf credentials are visible in the process list (ps aux) and shell history. Use a file path instead: --userinf=/etc/authwert/users.json')
            _p.userinf = json.loads(_p.userinf)
        elif os.path.isfile(_p.userinf):
            with open(_p.userinf) as f:
                _p.userinf = pb.Bag(json.loads(f.read()))
    if not _p.userinf:
        _p.userinf = {}

    # Warn about wildcard user
    if '*' in _p.userinf:
        Log('WARNING: Wildcard user (*) is configured. This allows a shared password for any username. Use with caution.')

    _p.run = True
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(run(_p))


if __name__ == '__main__':
    try:
        _p = pb.Bag({'threads': {}})
        main(_p)

    except KeyboardInterrupt:
        Log(" ~ keyboard ~ ")
    except Exception as e:
        Log(" ~ exception ~ ", e)
    finally:
        if _p.dbauth:
            _p.dbauth.close(_p)
        if _p.wsa:
            _p.wsa.stop()
            del _p.wsa
        Log("Bye...")
