#!/bin/bash

# Ensure required dependencies are installed
REQUIRED_CMDS=("curl" "jq" "uuidgen" "iptables" "iptables-save" "ip")
for cmd in "${REQUIRED_CMDS[@]}"; do
    if ! command -v "$cmd" &> /dev/null; then
        echo "Error: Required command '$cmd' is not installed. Install it and try again."
        exit 1
    fi
done

# Generate a GUID for Microsoft API request
CLIENT_REQUEST_ID=$(uuidgen)

# Fetch only the IPs for SMTP (port 25) from Office 365
O365_IPS=$(curl -s "https://endpoints.office.com/endpoints/worldwide?ClientRequestId=$CLIENT_REQUEST_ID" | jq -r '.[] | select(.serviceArea=="Exchange" and .tcpPorts=="25") | .ips[]' | grep -E "^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?$")

# If no IPs are found, exit
if [ -z "$O365_IPS" ]; then
    echo "Failed to retrieve Office 365 SMTP IPs. Exiting."
    exit 1
fi

# Flush old iptables rules
iptables -F
iptables -X

# Allow loopback traffic (localhost)
iptables -A INPUT -i lo -j ACCEPT

# Allow established and related connections
iptables -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT

# Detect and allow local subnet dynamically
LOCAL_SUBNET=$(ip -o -4 addr show | awk '$2 != "lo" {print $4}' | head -n1)
if [ -n "$LOCAL_SUBNET" ]; then
    echo "Allowing local subnet: $LOCAL_SUBNET"
    iptables -A INPUT -p tcp --dport 25 -s "$LOCAL_SUBNET" -j ACCEPT
else
    echo "Warning: Could not determine local subnet."
fi

# Allow SMTP traffic only from Office 365 Exchange Online servers
echo "Allowing Office 365 SMTP IPs:"
for ip in $O365_IPS; do
    echo "  - $ip"
    if ! iptables -C INPUT -p tcp --dport 25 -s "$ip" -j ACCEPT 2>/dev/null; then
        iptables -A INPUT -p tcp --dport 25 -s "$ip" -j ACCEPT
    else
        echo "(Rule for $ip already exists, skipping)"
    fi
done

# Block all other SMTP traffic
if ! iptables -C INPUT -p tcp --dport 25 -j DROP 2>/dev/null; then
    iptables -A INPUT -p tcp --dport 25 -j DROP
    echo "Added final rule to block all other SMTP traffic."
else
    echo "Final SMTP block rule already exists, skipping."
fi

# Save iptables rules for persistence
iptables-save > /etc/iptables/rules.v4
echo "$(date) - Updated Office 365 SMTP restrictions" >> /var/log/o365_smtp_restrictions.log
echo "Office 365 SMTP restrictions successfully applied!"