wol-sol-cli/index.py

#!/usr/bin/env python3
"""
Wake-on-LAN CLI tool with authentication support.
Sends a magic packet over UDP port 9999 with AUTH format.
"""

import socket
import argparse
import sys
from typing import Optional


def create_magic_packet(mac_address: str) -> bytes:
    """
    Creates a Wake-on-LAN magic packet.
    
    Args:
        mac_address: MAC address in format XX:XX:XX:XX:XX:XX or XX-XX-XX-XX-XX-XX
        
    Returns:
        Magic packet as bytes (6 bytes 0xFF + 16 repetitions of MAC address)
    """
    # Remove separators and convert to uppercase
    mac = mac_address.replace(':', '').replace('-', '').upper()
    
    # Validate MAC address format (should be 12 hex characters)
    if len(mac) != 12:
        raise ValueError(f"Invalid MAC address format: {mac_address}")
    
    try:
        # Convert MAC address to bytes
        mac_bytes = bytes.fromhex(mac)
    except ValueError as e:
        raise ValueError(f"Invalid MAC address format: {mac_address}") from e
    
    # Magic packet: 6 bytes of 0xFF followed by 16 repetitions of MAC address
    magic_packet = b'\xff' * 6 + mac_bytes * 16
    
    return magic_packet


def send_wol_packet(
    mac_address: str,
    username: str,
    password: str,
    host: str = "255.255.255.255",
    port: int = 9999
) -> bool:
    """
    Sends a Wake-on-LAN packet with authentication over UDP.
    
    Args:
        mac_address: MAC address of the target device
        username: Username for authentication
        password: Password for authentication
        host: Target host (default: broadcast address)
        port: Target port (default: 9999)
        
    Returns:
        True if packet was sent successfully, False otherwise
    """
    try:
        # Create magic packet
        magic_packet = create_magic_packet(mac_address)
        
        # Create authentication header
        auth_header = f"AUTH:{username}:{password}:".encode('utf-8')
        
        # Combine authentication header with magic packet
        packet = auth_header + magic_packet
        
        # Create UDP socket
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        
        # Enable broadcast
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        
        # Send packet
        sock.sendto(packet, (host, port))
        sock.close()
        
        return True
        
    except Exception as e:
        print(f"Error sending packet: {e}", file=sys.stderr)
        return False


def main():
    """Main entry point for the CLI application."""
    parser = argparse.ArgumentParser(
        description="Wake-on-LAN tool with authentication support",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    subparsers = parser.add_subparsers(dest='command', required=True, help='Available commands')

    # Wake command parser
    wake_parser = subparsers.add_parser(
        'wake',
        help='Send Wake-on-LAN packet (default port: 9)',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s wake -m AA:BB:CC:DD:EE:FF -u admin -p secret
  %(prog)s wake -m AA-BB-CC-DD-EE-FF -u user -p pass -H 192.168.1.100
  %(prog)s wake -m AA:BB:CC:DD:EE:FF
        """
    )

    wake_parser.add_argument(
        '-m', '--mac',
        required=True,
        help='MAC address of the target device (format: XX:XX:XX:XX:XX:XX or XX-XX-XX-XX-XX-XX)'
    )

    wake_parser.add_argument(
        '-u', '--username',
        help='Username for authentication (optional)'
    )

    wake_parser.add_argument(
        '-p', '--password',
        help='Password for authentication (optional)'
    )

    wake_parser.add_argument(
        '-H', '--host',
        default='255.255.255.255',
        help='Target host address (default: 255.255.255.255 for broadcast)'
    )

    wake_parser.add_argument(
        '-P', '--port',
        type=int,
        default=9,
        help='Target UDP port (default: 9)'
    )

    # Sleep command parser
    sleep_parser = subparsers.add_parser(
        'sleep',
        help='Send Sleep packet (default port: 9999)',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s sleep -m AA:BB:CC:DD:EE:FF -u admin -p secret
  %(prog)s sleep -m AA-BB-CC-DD-EE-FF -u user -p pass -H 192.168.1.100 -P 9999
        """
    )

    sleep_parser.add_argument(
        '-m', '--mac',
        required=True,
        help='MAC address of the target device (format: XX:XX:XX:XX:XX:XX or XX-XX-XX-XX-XX-XX)'
    )

    sleep_parser.add_argument(
        '-u', '--username',
        required=True,
        help='Username for authentication'
    )

    sleep_parser.add_argument(
        '-p', '--password',
        required=True,
        help='Password for authentication'
    )

    sleep_parser.add_argument(
        '-H', '--host',
        default='255.255.255.255',
        help='Target host address (default: 255.255.255.255 for broadcast)'
    )

    sleep_parser.add_argument(
        '-P', '--port',
        type=int,
        default=9999,
        help='Target UDP port (default: 9999)'
    )

    args = parser.parse_args()

    # For wake command, check if authentication is provided
    if args.command == 'wake':
        if (args.username and not args.password) or (args.password and not args.username):
            print("Error: Both username and password must be provided for authentication.", file=sys.stderr)
            return 1

    # Send the appropriate packet
    command_name = "Wake-on-LAN" if args.command == 'wake' else "Sleep"
    print(f"Sending {command_name} packet to {args.mac}...")
    print(f"Host: {args.host}:{args.port}")

    if args.username:
        print(f"Username: {args.username}")
    else:
        print("No authentication (anonymous)")

    success = send_wol_packet(
        mac_address=args.mac,
        username=args.username,
        password=args.password,
        host=args.host,
        port=args.port
    )

    if success:
        print("Packet sent successfully!")
        return 0
    else:
        print("Failed to send packet.", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())