#!/bin/sh

# =============================================================================
# OpenWRT Device Access Control Script
# =============================================================================
# This script manages device internet access through IP-based blocking
# 
# Usage:
#   ./device-control.sh block <IP> <NAME>      - Block device internet access
#   ./device-control.sh unblock <IP> <NAME>    - Allow device internet access
#   ./device-control.sh list                   - List all blocked devices
#   ./device-control.sh status <IP>            - Check if device is blocked
#   ./device-control.sh init                   - Initialize the blocked devices set
#
# Installation:
#   1. Copy to /root/device-control.sh on OpenWRT
#   2. chmod +x /root/device-control.sh
#   3. Run: ./device-control.sh init
#   4. Add to /etc/firewall.user for persistence
# =============================================================================

ACTION=$1
DEVICE_IP=$2
DEVICE_NAME=$3
IPSET_NAME="blocked_devices"
LOG_FILE="/var/log/device-control.log"

# Function to log actions
log_action() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') - $1" >> $LOG_FILE
}

# Function to initialize the ipset
init_ipset() {
    # Check if ipset exists
    if ! ipset list $IPSET_NAME >/dev/null 2>&1; then
        echo "Creating ipset: $IPSET_NAME"
        ipset create $IPSET_NAME hash:ip timeout 0 comment
        log_action "INIT: Created ipset $IPSET_NAME"
    else
        echo "ipset $IPSET_NAME already exists"
    fi
    
    # Check if firewall rule exists
    if ! iptables -C FORWARD -m set --match-set $IPSET_NAME src -j REJECT >/dev/null 2>&1; then
        echo "Adding firewall rule for blocking"
        iptables -I FORWARD -m set --match-set $IPSET_NAME src -j REJECT
        log_action "INIT: Added firewall rule"
    else
        echo "Firewall rule already exists"
    fi
    
    echo ""
    echo "Initialization complete!"
    echo "To make this persistent across reboots, add to /etc/firewall.user:"
    echo ""
    echo "ipset create $IPSET_NAME hash:ip timeout 0 comment -exist"
    echo "iptables -I FORWARD -m set --match-set $IPSET_NAME src -j REJECT"
}

# Function to validate IP address
validate_ip() {
    if [[ ! $1 =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
        echo "Error: Invalid IP address format"
        return 1
    fi
    return 0
}

# Main script logic
case $ACTION in
    init)
        init_ipset
        ;;
        
    block)
        if [ -z "$DEVICE_IP" ]; then
            echo "Error: IP address required"
            echo "Usage: $0 block <IP> [NAME]"
            exit 1
        fi
        
        validate_ip $DEVICE_IP || exit 1
        
        # Add to ipset with comment if name provided
        if [ -n "$DEVICE_NAME" ]; then
            ipset add $IPSET_NAME $DEVICE_IP comment "$DEVICE_NAME" -exist
            echo "✓ Blocked: $DEVICE_NAME ($DEVICE_IP)"
            log_action "BLOCK: $DEVICE_NAME ($DEVICE_IP)"
        else
            ipset add $IPSET_NAME $DEVICE_IP -exist
            echo "✓ Blocked: $DEVICE_IP"
            log_action "BLOCK: $DEVICE_IP"
        fi
        ;;
        
    unblock)
        if [ -z "$DEVICE_IP" ]; then
            echo "Error: IP address required"
            echo "Usage: $0 unblock <IP> [NAME]"
            exit 1
        fi
        
        validate_ip $DEVICE_IP || exit 1
        
        ipset del $IPSET_NAME $DEVICE_IP 2>/dev/null
        if [ $? -eq 0 ]; then
            if [ -n "$DEVICE_NAME" ]; then
                echo "✓ Unblocked: $DEVICE_NAME ($DEVICE_IP)"
                log_action "UNBLOCK: $DEVICE_NAME ($DEVICE_IP)"
            else
                echo "✓ Unblocked: $DEVICE_IP"
                log_action "UNBLOCK: $DEVICE_IP"
            fi
        else
            echo "✗ IP $DEVICE_IP was not in blocked list"
        fi
        ;;
        
    list)
        echo "========================================"
        echo "Currently Blocked Devices"
        echo "========================================"
        if ipset list $IPSET_NAME >/dev/null 2>&1; then
            ipset list $IPSET_NAME | grep -A 100 "Members:" | tail -n +2 | while read line; do
                if [ -n "$line" ]; then
                    echo "• $line"
                fi
            done
            
            # Count
            COUNT=$(ipset list $IPSET_NAME | grep -A 100 "Members:" | tail -n +2 | grep -c "")
            echo "========================================"
            echo "Total blocked: $COUNT device(s)"
        else
            echo "No blocked devices (ipset not initialized)"
            echo "Run: $0 init"
        fi
        ;;
        
    status)
        if [ -z "$DEVICE_IP" ]; then
            echo "Error: IP address required"
            echo "Usage: $0 status <IP>"
            exit 1
        fi
        
        validate_ip $DEVICE_IP || exit 1
        
        if ipset test $IPSET_NAME $DEVICE_IP >/dev/null 2>&1; then
            echo "🔴 $DEVICE_IP is BLOCKED"
            # Try to get comment
            COMMENT=$(ipset list $IPSET_NAME | grep "$DEVICE_IP" | grep -o 'comment ".*"' | sed 's/comment "\(.*\)"/\1/')
            if [ -n "$COMMENT" ]; then
                echo "   Device: $COMMENT"
            fi
            exit 0
        else
            echo "🟢 $DEVICE_IP is ALLOWED"
            exit 1
        fi
        ;;
        
    log)
        if [ -f "$LOG_FILE" ]; then
            echo "========================================"
            echo "Device Control Log"
            echo "========================================"
            tail -n 50 $LOG_FILE
        else
            echo "No log file found at $LOG_FILE"
        fi
        ;;
        
    clear)
        echo "WARNING: This will unblock ALL devices!"
        echo -n "Are you sure? (yes/no): "
        read CONFIRM
        if [ "$CONFIRM" = "yes" ]; then
            ipset flush $IPSET_NAME
            echo "✓ All devices unblocked"
            log_action "CLEAR: All devices unblocked"
        else
            echo "Cancelled"
        fi
        ;;
        
    help|--help|-h)
        echo "OpenWRT Device Access Control"
        echo ""
        echo "Usage:"
        echo "  $0 init                    - Initialize blocking system"
        echo "  $0 block <IP> [NAME]       - Block device internet access"
        echo "  $0 unblock <IP> [NAME]     - Allow device internet access"
        echo "  $0 list                    - List all blocked devices"
        echo "  $0 status <IP>             - Check if device is blocked"
        echo "  $0 log                     - View recent actions"
        echo "  $0 clear                   - Unblock all devices"
        echo "  $0 help                    - Show this help"
        echo ""
        echo "Examples:"
        echo "  $0 block 10.0.0.100 \"Kids Tablet\""
        echo "  $0 unblock 10.0.0.100"
        echo "  $0 status 10.0.0.100"
        echo ""
        ;;
        
    *)
        echo "Error: Unknown action '$ACTION'"
        echo "Run '$0 help' for usage information"
        exit 1
        ;;
esac

exit 0