pg-api 0.1.0

#!/usr/bin/env python3
"""
License Key Generator for PG-API
Generates offline license keys with HMAC validation
"""

import json
import hmac
import hashlib
import base64
import secrets
import argparse
from datetime import datetime, timedelta
from typing import Dict, Optional

# Secret key for HMAC (in production, this should be kept secure)
# This same key must be compiled into the pg-api binary
SECRET_KEY = "pg-api-license-secret-2024-production"

class LicenseType:
    TRIAL = "trial"
    STANDARD = "standard"
    ENTERPRISE = "enterprise"

class LicenseGenerator:
    def __init__(self, secret_key: str = SECRET_KEY):
        self.secret_key = secret_key.encode('utf-8')
    
    def generate_license(
        self,
        customer_name: str,
        license_type: str = LicenseType.STANDARD,
        days_valid: int = 365,
        max_connections: Optional[int] = None
    ) -> Dict[str, str]:
        """Generate a license key with metadata"""
        
        # Set defaults based on license type
        if max_connections is None:
            max_connections = {
                LicenseType.TRIAL: 10,
                LicenseType.STANDARD: 100,
                LicenseType.ENTERPRISE: 1000
            }.get(license_type, 100)
        
        # Generate license ID
        license_id = secrets.token_hex(8).upper()
        
        # Calculate expiration
        issued_at = datetime.utcnow()
        expires_at = issued_at + timedelta(days=days_valid)
        
        # Create license payload
        payload = {
            "id": license_id,
            "customer": customer_name,
            "type": license_type,
            "issued": issued_at.isoformat(),
            "expires": expires_at.isoformat(),
            "max_connections": max_connections,
            "version": "1"
        }
        
        # Create license data string (compact format for key)
        license_data = f"{license_id}|{license_type}|{expires_at.strftime('%Y%m%d')}|{max_connections}"
        
        # Generate HMAC signature
        signature = hmac.new(
            self.secret_key,
            license_data.encode('utf-8'),
            hashlib.sha256
        ).digest()
        
        # Encode signature to base64
        signature_b64 = base64.urlsafe_b64encode(signature).decode('utf-8').rstrip('=')
        
        # Create the license key
        license_key = f"PG-{license_type.upper()[:3]}-{license_id}-{signature_b64[:16]}"
        
        # Add key to payload
        payload["key"] = license_key
        
        return payload
    
    def validate_license(self, license_key: str) -> bool:
        """Validate a license key (for testing)"""
        try:
            # Parse license key
            parts = license_key.split('-')
            if len(parts) != 4 or parts[0] != "PG":
                return False
            
            license_type_abbr = parts[1]
            license_id = parts[2]
            provided_signature = parts[3]
            
            # This is a simplified validation
            # The actual validation would need to decode and verify the HMAC
            return len(license_id) == 16 and len(provided_signature) >= 10
            
        except Exception:
            return False
    
    def save_license_file(self, license_data: Dict, filename: str):
        """Save license to a JSON file"""
        with open(filename, 'w') as f:
            json.dump(license_data, f, indent=2)
        print(f"License saved to: {filename}")
    
    def save_license_info(self, license_data: Dict, filename: str):
        """Save human-readable license information"""
        with open(filename, 'w') as f:
            f.write("=" * 60 + "\n")
            f.write("PG-API LICENSE INFORMATION\n")
            f.write("=" * 60 + "\n\n")
            
            f.write(f"License ID:       {license_data['id']}\n")
            f.write(f"Customer:         {license_data['customer']}\n")
            f.write(f"License Type:     {license_data['type'].upper()}\n")
            f.write(f"Issue Date:       {license_data['issued']}\n")
            f.write(f"Expiration Date:  {license_data['expires']}\n")
            f.write(f"Max Connections:  {license_data['max_connections']}\n")
            f.write(f"License Version:  {license_data['version']}\n")
            f.write("\n")
            f.write("-" * 60 + "\n")
            f.write("LICENSE KEY:\n")
            f.write("-" * 60 + "\n")
            f.write(f"{license_data['key']}\n")
            f.write("-" * 60 + "\n\n")
            
            f.write("INSTALLATION:\n")
            f.write("Add this key to your .env file:\n")
            f.write(f"LICENSE_KEY={license_data['key']}\n\n")
            
            f.write("Or use during installation:\n")
            f.write(f"sudo python3 setup.py --license-key \"{license_data['key']}\"\n")
            f.write("\n" + "=" * 60 + "\n")
        
        print(f"License info saved to: {filename}")

def main():
    parser = argparse.ArgumentParser(
        description='Generate license keys for PG-API',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    parser.add_argument(
        'customer',
        help='Customer name or organization'
    )
    
    parser.add_argument(
        '--type',
        choices=[LicenseType.TRIAL, LicenseType.STANDARD, LicenseType.ENTERPRISE],
        default=LicenseType.STANDARD,
        help='License type (default: standard)'
    )
    
    parser.add_argument(
        '--days',
        type=int,
        default=365,
        help='Days until expiration (default: 365)'
    )
    
    parser.add_argument(
        '--max-connections',
        type=int,
        help='Maximum database connections (default: based on type)'
    )
    
    parser.add_argument(
        '--output',
        default='license',
        help='Output filename prefix (default: license)'
    )
    
    parser.add_argument(
        '--validate',
        help='Validate an existing license key'
    )
    
    args = parser.parse_args()
    
    generator = LicenseGenerator()
    
    # Validate mode
    if args.validate:
        is_valid = generator.validate_license(args.validate)
        if is_valid:
            print(f"✓ License key appears valid: {args.validate}")
        else:
            print(f"✗ Invalid license key: {args.validate}")
        return
    
    # Generate mode
    print("Generating PG-API License...")
    print("-" * 40)
    
    license_data = generator.generate_license(
        customer_name=args.customer,
        license_type=args.type,
        days_valid=args.days,
        max_connections=args.max_connections
    )
    
    # Display license info
    print(f"Customer:        {license_data['customer']}")
    print(f"License Type:    {license_data['type'].upper()}")
    print(f"License ID:      {license_data['id']}")
    print(f"Valid Until:     {license_data['expires']}")
    print(f"Max Connections: {license_data['max_connections']}")
    print("-" * 40)
    print(f"License Key:     {license_data['key']}")
    print("-" * 40)
    
    # Save files
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    json_file = f"{args.output}_{timestamp}.json"
    info_file = f"{args.output}_{timestamp}.txt"
    
    generator.save_license_file(license_data, json_file)
    generator.save_license_info(license_data, info_file)
    
    print("\n✓ License generation complete!")
    print(f"\nTo use this license, add to your .env file:")
    print(f"LICENSE_KEY={license_data['key']}")

if __name__ == '__main__':
    main()