secrettext/security.py

import time
import hashlib
import json
from collections import defaultdict
from flask import request, render_template
import redis
import os
import database

redis_client = redis.Redis(
    host=os.environ.get('REDIS_HOST', 'localhost'),
    port=int(os.environ.get('REDIS_PORT', 6379)),
    db=int(os.environ.get('REDIS_DB', 0)),
    password=os.environ.get('REDIS_PASSWORD', None),
    decode_responses=True
)

BLACKLIST_PREFIX = 'blacklist:'
ATTEMPTS_PREFIX = 'attempts:'
BLOCKED_PREFIX = 'blocked:'

SECURITY_CONFIG = {
    'max_attempts': 5,
    'block_duration': 300,
    'block_duration_increment': 60,
    'max_block_duration': 3600,
    'attempt_window': 60,
    'permanent_blacklist': False,
    'permanent_after_blocks': 3,
}

admin_attempts = defaultdict(list)


def get_client_identifier():
    ip = request.remote_addr
    user_agent = request.headers.get('User-Agent', 'unknown')[:100]
    identifier = hashlib.sha256(f"{ip}:{user_agent}".encode()).hexdigest()
    return identifier, ip


def get_ip_block_info(ip: str) -> dict:
    if redis_client.exists(f"{BLACKLIST_PREFIX}{ip}"):
        return {'blocked': True, 'permanent': True, 'reason': 'Постоянная блокировка', 'until': None}

    block_key = f"{BLOCKED_PREFIX}{ip}"
    ttl = redis_client.ttl(block_key)
    if ttl > 0:
        return {
            'blocked': True,
            'permanent': False,
            'reason': 'Слишком много попыток',
            'until': time.time() + ttl,
            'remaining_seconds': ttl,
            'remaining_minutes': ttl // 60
        }
    return {'blocked': False}


def get_attempts_count(ip: str) -> int:
    attempts_key = f"{ATTEMPTS_PREFIX}{ip}"
    attempts_data = redis_client.get(attempts_key)
    if not attempts_data:
        return 0
    try:
        attempts = json.loads(attempts_data)
        current_time = time.time()
        valid_attempts = [t for t in attempts if current_time - t < SECURITY_CONFIG['attempt_window']]
        return len(valid_attempts)
    except:
        return 0


def block_ip(ip: str):
    block_count_key = f"{BLOCKED_PREFIX}{ip}:count"
    block_count = int(redis_client.get(block_count_key) or 0)

    block_duration = min(
        SECURITY_CONFIG['block_duration'] + (block_count * SECURITY_CONFIG['block_duration_increment']),
        SECURITY_CONFIG['max_block_duration']
    )

    redis_client.setex(f"{BLOCKED_PREFIX}{ip}", block_duration, '1')
    redis_client.setex(block_count_key, 86400 * 30, str(block_count + 1))
    redis_client.delete(f"{ATTEMPTS_PREFIX}{ip}")

    database.log_security_event(
        ip=ip,
        event_type='temporary_block',
        details=f'Блокировка на {block_duration} секунд (попытка #{block_count + 1})'
    )


def record_failed_attempt(identifier: str):
    _, ip = get_client_identifier()
    attempts_key = f"{ATTEMPTS_PREFIX}{ip}"
    current_time = time.time()

    attempts_data = redis_client.get(attempts_key)
    attempts = []
    if attempts_data:
        try:
            attempts = json.loads(attempts_data)
            attempts = [t for t in attempts if current_time - t < SECURITY_CONFIG['attempt_window']]
        except:
            attempts = []

    attempts.append(current_time)
    redis_client.setex(attempts_key, SECURITY_CONFIG['attempt_window'] * 2, json.dumps(attempts))

    if len(attempts) >= SECURITY_CONFIG['max_attempts']:
        block_ip(ip)


def clear_failed_attempts(identifier: str):
    _, ip = get_client_identifier()
    redis_client.delete(f"{ATTEMPTS_PREFIX}{ip}")


def get_blocked_ips() -> list:
    result = []

    cursor = 0
    try:
        while True:
            cursor, keys = redis_client.scan(
                cursor,
                match=f"{BLOCKED_PREFIX}*",
                count=100
            )

            for key in keys:
                if key.endswith(':count'):
                    continue
                ip = key.replace(BLOCKED_PREFIX, '')
                ttl = redis_client.ttl(key)
                if ttl > 0:
                    count_key = f"{BLOCKED_PREFIX}{ip}:count"
                    count = int(redis_client.get(count_key) or 1)
                    result.append({
                        'ip': ip,
                        'permanent': False,
                        'remaining_seconds': ttl,
                        'remaining_minutes': ttl // 60,
                        'block_count': count
                    })

            if cursor == 0:
                break
    except Exception as e:
        print(f"Redis scan error in get_blocked_ips: {e}")
        return result

    cursor = 0
    try:
        while True:
            cursor, keys = redis_client.scan(
                cursor,
                match=f"{BLACKLIST_PREFIX}*",
                count=100
            )

            for key in keys:
                ip = key.replace(BLACKLIST_PREFIX, '')
                result.append({'ip': ip, 'permanent': True, 'block_count': '∞'})

            if cursor == 0:
                break
    except Exception as e:
        print(f"Redis scan error in get_blocked_ips (blacklist): {e}")
        return result

    return result


def unblock_ip(ip: str):
    redis_client.delete(f"{BLACKLIST_PREFIX}{ip}")
    redis_client.delete(f"{BLOCKED_PREFIX}{ip}")
    redis_client.delete(f"{ATTEMPTS_PREFIX}{ip}")
    redis_client.delete(f"{BLOCKED_PREFIX}{ip}:count")
    return True


def check_admin_rate_limit(ip: str) -> bool:
    current_time = time.time()
    admin_attempts[ip] = [t for t in admin_attempts[ip] if current_time - t < 300]
    if len(admin_attempts[ip]) >= 5:
        return True
    admin_attempts[ip].append(current_time)
    return False