Basic Diameter Examples¶

from diameter.client import DiameterClient
from diameter.avp import AVP
import asyncio
import uuid

class BasicDiameterClient:
    def __init__(self, host, port, realm):
        self.host = host
        self.port = port
        self.realm = realm
        self.client = DiameterClient()

    async def connect(self):
        """Establish connection and perform CER/CEA"""
        await self.client.connect(self.host, self.port)
        await self.capability_exchange()

    async def capability_exchange(self):
        """Perform Capabilities-Exchange"""
        avps = [
            AVP('Origin-Host', self.host),
            AVP('Origin-Realm', self.realm),
            AVP('Host-IP-Address', self.host),
            AVP('Vendor-Id', 0),
            AVP('Product-Name', 'Example Client'),
            AVP('Auth-Application-Id', 4)  # Credit Control
        ]

        request = self.client.create_request('CER', avps)
        response = await self.client.send_request(request)
        return response.result_code == 2001

async def main():
    client = BasicDiameterClient(
        host='client.example.com',
        port=3868,
        realm='example.com'
    )

    try:
        await client.connect()
        print("Connected successfully")

    except Exception as e:
        print(f"Connection failed: {e}")

if __name__ == '__main__':
    asyncio.run(main())

async def send_auth_request(client, username):
    """Send AA-Request for user authentication"""
    session_id = str(uuid.uuid4())

    avps = [
        AVP('Session-Id', session_id),
        AVP('Origin-Host', client.host),
        AVP('Origin-Realm', client.realm),
        AVP('Destination-Realm', 'auth.example.com'),
        AVP('Auth-Application-Id', 1),
        AVP('User-Name', username),
        AVP('Auth-Request-Type', 1)  # AUTHENTICATE_ONLY
    ]

    request = client.create_request('AAR', avps)
    return await client.send_request(request)

# Usage example
async def authenticate_user(client, username):
    try:
        response = await send_auth_request(client, username)

        if response.result_code == 2001:
            print(f"Authentication successful for {username}")
            return True
        else:
            print(f"Authentication failed: {response.error_message}")
            return False

    except Exception as e:
        print(f"Error during authentication: {e}")
        return False

async def send_credit_control_request(client, subscription_id, requested_units):
    """Send Credit-Control-Request for real-time charging"""
    session_id = str(uuid.uuid4())

    avps = [
        AVP('Session-Id', session_id),
        AVP('Origin-Host', client.host),
        AVP('Origin-Realm', client.realm),
        AVP('Destination-Realm', 'charging.example.com'),
        AVP('Auth-Application-Id', 4),
        AVP('CC-Request-Type', 1),  # INITIAL_REQUEST
        AVP('CC-Request-Number', 0),
        AVP('Subscription-Id', [
            AVP('Subscription-Id-Type', 0),  # END_USER_E164
            AVP('Subscription-Id-Data', subscription_id)
        ]),
        AVP('Requested-Service-Unit', [
            AVP('CC-Time', requested_units)
        ])
    ]

    request = client.create_request('CCR', avps)
    return await client.send_request(request)

# Usage example
async def request_credit(client, subscription_id):
    try:
        response = await send_credit_control_request(
            client,
            subscription_id,
            requested_units=600  # 10 minutes
        )

        if response.result_code == 2001:
            granted_units = response.get_avp('Granted-Service-Unit')
            print(f"Granted {granted_units} units")
            return granted_units
        else:
            print(f"Credit request failed: {response.error_message}")
            return 0

    except Exception as e:
        print(f"Error during credit request: {e}")
        return 0

class DiameterError(Exception):
    def __init__(self, result_code, message=None):
        self.result_code = result_code
        self.message = message or self._get_error_message()
        super().__init__(self.message)

    def _get_error_message(self):
        error_messages = {
            3001: "DIAMETER_COMMAND_UNSUPPORTED",
            3002: "DIAMETER_UNABLE_TO_DELIVER",
            3004: "DIAMETER_TOO_BUSY",
            4001: "DIAMETER_AUTHENTICATION_REJECTED",
            4010: "DIAMETER_END_USER_SERVICE_DENIED",
            4012: "DIAMETER_CREDIT_LIMIT_REACHED",
            5001: "DIAMETER_AVP_UNSUPPORTED",
            5004: "DIAMETER_INVALID_AVP_VALUE",
            5012: "DIAMETER_UNABLE_TO_COMPLY"
        }
        return error_messages.get(
            self.result_code,
            f"Unknown error: {self.result_code}"
        )

async def send_request_with_retry(client, request, max_retries=3):
    """Send request with retry mechanism"""
    for attempt in range(max_retries):
        try:
            response = await client.send_request(request)

            if response.result_code == 2001:
                return response

            if response.result_code >= 3000:
                raise DiameterError(response.result_code)

        except Exception as e:
            if attempt == max_retries - 1:
                raise

            await asyncio.sleep(2 ** attempt)  # Exponential backoff
            continue

# config.yaml
diameter:
  client:
    host: client.example.com
    port: 3868
    realm: example.com

  connection:
    retry_count: 3
    timeout: 5
    keepalive: 30

  security:
    tls_enabled: true
    cert_file: /etc/diameter/client.pem
    key_file: /etc/diameter/client.key

  logging:
    level: INFO
    file: /var/log/diameter.log

import yaml

def load_config(config_file):
    """Load diameter configuration from YAML file"""
    with open(config_file, 'r') as f:
        return yaml.safe_load(f)

def create_client_from_config(config_file):
    """Create Diameter client from configuration"""
    config = load_config(config_file)

    client = BasicDiameterClient(
        host=config['diameter']['client']['host'],
        port=config['diameter']['client']['port'],
        realm=config['diameter']['client']['realm']
    )

    if config['diameter']['security']['tls_enabled']:
        client.set_tls_config(
            cert_file=config['diameter']['security']['cert_file'],
            key_file=config['diameter']['security']['key_file']
        )

    return client