JWT Service-to-Service Authentication Mock in Python

import hashlib
import hmac
import base64
import json
import time


class JWTMock:
    """Minimal JWT service-to-service mock using HS256."""
    
    def __init__(self, secret):
        self.secret = secret.encode()
    
    @staticmethod
    def _b64url_encode(data):
        return base64.urlsafe_b64encode(data).rstrip(b'=').decode()
    
    @staticmethod
    def _b64url_decode(data):
        padding = '=' * (-len(data) % 4)
        return base64.urlsafe_b64decode(data + padding)
    
    def create_token(self, service_name, audience, ttl_seconds=3600):
        header = {"alg": "HS256", "typ": "JWT"}
        payload = {
            "iss": service_name,
            "sub": service_name,
            "aud": audience,
            "iat": int(time.time()),
            "exp": int(time.time()) + ttl_seconds
        }
        header_b64 = self._b64url_encode(json.dumps(header).encode())
        payload_b64 = self._b64url_encode(json.dumps(payload).encode())
        signing_input = f"{header_b64}.{payload_b64}".encode()
        signature = hmac.new(self.secret, signing_input, hashlib.sha256).digest()
        return f"{header_b64}.{payload_b64}.{self._b64url_encode(signature)}"
    
    def verify_token(self, token, expected_audience):
        try:
            header_b64, payload_b64, signature_b64 = token.split('.')
            signing_input = f"{header_b64}.{payload_b64}".encode()
            expected_sig = hmac.new(self.secret, signing_input, hashlib.sha256).digest()
            if not hmac.compare_digest(expected_sig, self._b64url_decode(signature_b64)):
                return False
            payload = json.loads(self._b64url_decode(payload_b64))
            return (payload.get("aud") == expected_audience 
                    and payload.get("exp", 0) > time.time())
        except Exception:
            return False


if __name__ == "__main__":
    jwt = JWTMock("shared-secret-key")
    auth_service = jwt.create_token("auth-service", "payment-service")
    print("Token:", auth_service)
    print("Verified by payment-service:", jwt.verify_token(auth_service, "payment-service"))
    print("Rejected by wrong audience:", jwt.verify_token(auth_service, "other-service"))