url safe base64 encode/decode

b64u.py

#!/usr/bin/env python3
import base64
import sys

def encode_bytes(data):
    return base64.urlsafe_b64encode(data).decode().rstrip('=')

def decode_bytes(data):
    padding = 4 - len(data) % 4
    if padding != 4:
        data += '=' * padding
    return base64.urlsafe_b64decode(data.encode())

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} encode|decode [file|-]")
        print("  - for stdin, omit file for argument input")
        sys.exit(1)
    
    action = sys.argv[1]
    
    # Determine input source
    if len(sys.argv) >= 3:
        filename = sys.argv[2]
        if filename == "-":
            # Read binary from stdin
            data = sys.stdin.buffer.read()
        else:
            # Read binary from file
            try:
                with open(filename, 'rb') as f:
                    data = f.read()
            except FileNotFoundError:
                print(f"Error: File not found: {filename}", file=sys.stderr)
                sys.exit(1)
            except Exception as e:
                print(f"Error reading file: {e}", file=sys.stderr)
                sys.exit(1)
    else:
        # Read argument from command line (string)
        if not sys.stdin.isatty():
            # Data is piped to stdin as text
            data = sys.stdin.read().rstrip('\n').encode()
        else:
            # No data provided
            print("Error: No input data provided", file=sys.stderr)
            sys.exit(1)
    
    try:
        if action == "encode":
            result = encode_bytes(data)
            sys.stdout.write(result)
        elif action == "decode":
            result = decode_bytes(data.decode() if isinstance(data, bytes) else data)
            sys.stdout.buffer.write(result)
        else:
            print(f"Unknown action: {action}", file=sys.stderr)
            sys.exit(1)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)