Tools API¶

def generate_key(
    key_type: str, bits: Optional[int] = None, comment: Optional[str] = None
) -> tuple[Any, Any]:
    """Generate a new SSH key pair."""
    from ..crypto.pkey import PKey

    key: PKey
    if key_type == "ed25519":
        key = Ed25519Key.generate()
    elif key_type == "ecdsa":
        key = ECDSAKey.generate(bits=bits or 256)
    elif key_type == "rsa":
        key_size = bits or 2048
        if key_size < 2048:
            raise ValueError("RSA key size must be at least 2048 bits")
        key = RSAKey.generate(bits=key_size)
    else:
        raise ValueError(f"Unsupported key type: {key_type}")

    return key, key.get_public_key()

def main() -> None:
    """Main entry point for ssh-keygen tool."""
    parser = argparse.ArgumentParser(
        description="Generate SSH key pairs",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  spindlex-keygen -t ed25519 -f ~/.ssh/id_ed25519
  spindlex-keygen -t rsa -b 4096 -f ~/.ssh/id_rsa -C "user@example.com"
  spindlex-keygen -t ecdsa -b 384 -f ~/.ssh/id_ecdsa
        """,
    )

    parser.add_argument(
        "-t",
        "--type",
        choices=["ed25519", "ecdsa", "rsa"],
        default="ed25519",
        help="Key type to generate (default: ed25519)",
    )

    parser.add_argument(
        "-b",
        "--bits",
        type=int,
        help="Number of bits for RSA keys (min 2048) or ECDSA keys (256, 384, 521)",
    )

    parser.add_argument(
        "-f", "--filename", required=True, help="Output filename for the private key"
    )

    parser.add_argument("-C", "--comment", help="Comment to add to the public key")

    parser.add_argument(
        "--overwrite", action="store_true", help="Overwrite existing key files"
    )

    args = parser.parse_args()

    # Check if files already exist
    private_path = Path(args.filename)
    public_path = Path(f"{args.filename}.pub")

    if not args.overwrite:
        if private_path.exists():
            print(f"Error: Private key file already exists: {private_path}")
            sys.exit(1)
        if public_path.exists():
            print(f"Error: Public key file already exists: {public_path}")
            sys.exit(1)

    try:
        # Generate key pair
        print(f"Generating {args.type} key pair...")
        private_key, public_key = generate_key(args.type, args.bits, args.comment)

        # Save key pair
        save_key_pair(private_key, public_key, args.filename, args.comment)

        # Show fingerprint
        fingerprint = public_key.get_fingerprint()
        print(f"Key fingerprint: {fingerprint}")
        print("Generated with Spindle SSH key generator")

    except Exception as e:
        print(f"Error generating key: {e}")
        sys.exit(1)

def save_key_pair(
    private_key: Any, public_key: Any, filename: str, comment: Optional[str] = None
) -> None:
    """Save the key pair to files."""
    private_path = Path(filename)
    public_path = Path(f"{filename}.pub")

    # Save private key
    private_key.save_to_file(str(private_path))
    private_path.chmod(0o600)  # Secure permissions

    # Save public key
    public_key_str = public_key.get_openssh_string()
    if comment:
        public_key_str += f" {comment}"

    public_path.write_text(public_key_str + "\n")
    public_path.chmod(0o644)

    print(f"Private key saved to: {private_path}")
    print(f"Public key saved to: {public_path}")

def benchmark_command_execution(
    client: SSHClient, command: str, iterations: int = 10
) -> dict[str, Any]:
    """Benchmark command execution."""
    times = []

    for _i in range(iterations):
        start_time = time.time()

        stdin, stdout, stderr = client.exec_command(command)
        stdout.read()  # Read all output
        stderr.read()  # Read all errors

        exec_time = time.time() - start_time
        times.append(exec_time)

    return {
        "operation": f"exec_command: {command}",
        "iterations": iterations,
        "times": times,
        "mean": statistics.mean(times),
        "median": statistics.median(times),
        "stdev": statistics.stdev(times) if len(times) > 1 else 0,
        "min": min(times),
        "max": max(times),
    }

def benchmark_connection(
    hostname: str,
    username: str,
    password: Optional[str] = None,
    key_filename: Optional[str] = None,
    iterations: int = 10,
) -> dict[str, Any]:
    """Benchmark SSH connection establishment."""
    times = []

    for _i in range(iterations):
        start_time = time.time()

        client = SSHClient()
        try:
            client.connect(
                hostname=hostname,
                username=username,
                password=password,
                key_filename=key_filename,
                timeout=30,
            )
            connect_time = time.time() - start_time
            times.append(connect_time)
        finally:
            client.close()

    return {
        "operation": "connection",
        "iterations": iterations,
        "times": times,
        "mean": statistics.mean(times),
        "median": statistics.median(times),
        "stdev": statistics.stdev(times) if len(times) > 1 else 0,
        "min": min(times),
        "max": max(times),
    }

def benchmark_crypto_operations(iterations: int = 1000) -> list[dict[str, Any]]:
    """Benchmark cryptographic operations."""
    get_crypto_backend()
    results = []

    # Benchmark key generation
    for key_type in ["ed25519", "ecdsa", "rsa"]:
        times = []

        for _i in range(min(iterations, 10)):  # Limit key gen iterations
            start_time = time.time()

            if key_type == "ed25519":
                from ..crypto.pkey import Ed25519Key

                Ed25519Key.generate()
            elif key_type == "ecdsa":
                from ..crypto.pkey import ECDSAKey

                ECDSAKey.generate()
            elif key_type == "rsa":
                from ..crypto.pkey import RSAKey

                RSAKey.generate(bits=2048)

            gen_time = time.time() - start_time
            times.append(gen_time)

        results.append(
            {
                "operation": f"{key_type}_key_generation",
                "iterations": len(times),
                "times": times,
                "mean": statistics.mean(times),
                "median": statistics.median(times),
                "stdev": statistics.stdev(times) if len(times) > 1 else 0,
                "min": min(times),
                "max": max(times),
            }
        )

    return results

def main() -> int:
    """Main entry point for ssh-benchmark tool."""
    parser = argparse.ArgumentParser(
        description="Benchmark SSH operations",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  spindlex-benchmark --crypto-only
  spindlex-benchmark -H example.com -u user -p password
  spindlex-benchmark -H example.com -u user -k ~/.ssh/id_rsa --iterations 20
        """,
    )

    parser.add_argument("-H", "--hostname", help="SSH server hostname")

    parser.add_argument("-u", "--username", help="SSH username")

    parser.add_argument(
        "-p", "--password", action="store_true", help="Prompt for SSH password"
    )

    parser.add_argument("-k", "--key-filename", help="SSH private key file")

    parser.add_argument(
        "-i",
        "--iterations",
        type=int,
        default=10,
        help="Number of iterations for each benchmark (default: 10)",
    )

    parser.add_argument(
        "--crypto-only", action="store_true", help="Only run cryptographic benchmarks"
    )

    parser.add_argument(
        "--commands",
        nargs="+",
        default=["echo hello", "ls -la", "uname -a"],
        help="Commands to benchmark (default: echo hello, ls -la, uname -a)",
    )

    args = parser.parse_args()

    results = []

    # Always run crypto benchmarks
    print("Running cryptographic benchmarks...")
    crypto_results = benchmark_crypto_operations(args.iterations)
    results.extend(crypto_results)

    if not args.crypto_only:
        if not args.hostname or not args.username:
            print("Error: hostname and username required for SSH benchmarks")
            parser.print_help()
            return 1

        if not args.password and not args.key_filename:
            print("Error: either password or key filename required")
            parser.print_help()
            return 1

        password = getpass.getpass("SSH password: ") if args.password else None

        # Benchmark connection
        print(f"Benchmarking connections to {args.hostname}...")
        conn_result = benchmark_connection(
            args.hostname,
            args.username,
            password,
            args.key_filename,
            args.iterations,
        )
        results.append(conn_result)

        # Benchmark commands
        client = SSHClient()
        try:
            client.connect(
                hostname=args.hostname,
                username=args.username,
                password=password,
                key_filename=args.key_filename,
            )

            for command in args.commands:
                print(f"Benchmarking command: {command}")
                cmd_result = benchmark_command_execution(
                    client, command, args.iterations
                )
                results.append(cmd_result)

        finally:
            client.close()

    # Print results
    print_benchmark_results(results)

    return 0

def print_benchmark_results(results: list[dict[str, Any]]) -> None:
    """Print benchmark results in a formatted table."""
    print("\nBenchmark Results")
    print("=" * 80)
    print(
        f"{'Operation':<30} {'Iterations':<10} {'Mean (s)':<12} {'Median (s)':<12} {'StdDev (s)':<12}"
    )
    print("-" * 80)

    for result in results:
        print(
            f"{result['operation']:<30} {result['iterations']:<10} "
            f"{result['mean']:<12.4f} {result['median']:<12.4f} {result['stdev']:<12.4f}"
        )