#!/usr/bin/env python3
"""
Файрволл на Python с использованием iptables
"""

import subprocess
import os
import sys


class IptablesFirewall:
    def __init__(self):
        self.chain_name = "PYTHON_FIREWALL"

    def run_command(self, cmd):
        """Выполнение команды iptables"""
        try:
            subprocess.run(cmd, shell=True, check=True)
            return True
        except subprocess.CalledProcessError as e:
            print(f"Error executing command: {cmd}")
            return False

    def initialize(self):
        """Инициализация цепочки файрволла"""
        # Создание пользовательской цепочки
        self.run_command(f"iptables -N {self.chain_name}")
        # Переход к нашей цепочке из INPUT
        self.run_command(f"iptables -I INPUT -j {self.chain_name}")

    def add_rule(self, protocol=None, source_ip=None, dest_port=None, action="DROP"):
        """Добавление правила в iptables"""
        cmd = f"iptables -A {self.chain_name}"

        if protocol:
            cmd += f" -p {protocol}"
        if source_ip:
            cmd += f" -s {source_ip}"
        if dest_port:
            cmd += f" --dport {dest_port}"

        cmd += f" -j {action}"

        return self.run_command(cmd)

    def block_ip(self, ip_address):
        """Блокировка IP-адреса"""
        return self.add_rule(source_ip=ip_address, action="DROP")

    def allow_port(self, port, protocol="tcp"):
        """Разрешение порта"""
        return self.add_rule(protocol=protocol, dest_port=port, action="ACCEPT")

    def block_port(self, port, protocol="tcp"):
        """Блокировка порта"""
        return self.add_rule(protocol=protocol, dest_port=port, action="DROP")

    def cleanup(self):
        """Очистка правил"""
        self.run_command(f"iptables -D INPUT -j {self.chain_name}")
        self.run_command(f"iptables -F {self.chain_name}")
        self.run_command(f"iptables -X {self.chain_name}")


# Пример использования
if __name__ == "__main__":
    if os.geteuid() != 0:
        print("Requires root privileges")
        sys.exit(1)

    fw = IptablesFirewall()

    try:
        fw.initialize()
        fw.block_port(23)  # Блокировка Telnet
        fw.allow_port(22)  # Разрешение SSH
        fw.allow_port(80)
        fw.allow_port(443)  # Разрешение HTTPS
        fw.block_ip("10.0.12.206")

        print("Firewall rules applied. Press Enter to cleanup...")
        input()

    finally:
        fw.cleanup()
        print("Firewall rules cleaned up")