controles/scripts/create_admin.py

import os
import pyotp
from pathlib import Path
from functions.database import Usuario, Role, get_db_session
from services.otp_service import generate_qr_code

ADMIN_USERNAME = "admin"
ADMIN_PASSWORD = "admin123"
ADMIN_ROLE = Role.SECRETARIO_GERAL


def salvar_qr_code(user):
    """
    Gera o QR code para um usuário específico
    
    Args:
        user: Instância do modelo Usuario
        
    Returns:
        tuple: (caminho do arquivo, URI do OTP)
    """
    # Tentar diferentes caminhos para salvar o QR code
    qr_paths = [
        Path('/tmp/admin_qr.png'),   # Diretório temporário do sistema
        Path('/data/admin_qr.png'),  # Diretório de dados do container
        Path('admin_qr.png')         # Diretório atual (fallback fora do container)
    ]
    
    # Tentar salvar em diferentes locais
    qr_saved = False
    saved_path = None
    
    img = generate_qr_code(user)  # Gera o QR code para o usuário

    for qr_path in qr_paths:
        try:
            # Tentar salvar o arquivo
            img.save(qr_path)
            qr_saved = True
            saved_path = qr_path
            break
        except Exception as e:
            print(f"Não foi possível salvar o QR code em {qr_path}: {e}")
            continue
    
    if not qr_saved:
        print("AVISO: Não foi possível salvar o QR code em nenhum local")
        print("O QR code pode ser gerado manualmente usando o URI OTP")
        saved_path = None
    
    return saved_path


def _ensure_admin_role(db, admin_user, role):
    admin_role = db.query(Role).filter_by(nome=role).first()
    if admin_role is None:
        admin_role = Role(nome=role, nivel=Role.SECRETARIO_GERAL)
        db.add(admin_role)
        db.flush()

    if admin_role not in admin_user.roles:
        admin_user.roles.append(admin_role)
        db.flush()


def _ensure_admin_otp(db, admin_user):
    if admin_user.otp_secret:
        return False
    secret = (os.environ.get('ADMIN_OTP_SECRET') or "").strip()
    admin_user.otp_secret = secret or admin_user.generate_otp_secret()
    db.flush()
    return True


def create_admin(username=ADMIN_USERNAME, password=ADMIN_PASSWORD, role=ADMIN_ROLE, save_qr=True):
    """Limpa e cria o usuário admin"""
    db = get_db_session()
    try:
        # Verificar se já existe um usuário admin
        admin_user = db.query(Usuario).filter_by(username=username).first()
        if admin_user is not None:
            db.delete(admin_user)
            db.flush()

        print("\n=== Criando Novo Usuário Admin ===")
        admin_user = Usuario(
            username=username,
            email="admin@example.com",
            is_admin=True
        )
        admin_user.set_password(password)
        _ensure_admin_otp(db, admin_user)
        _ensure_admin_role(db, admin_user, role)
        
        db.add(admin_user)
        db.commit()
        
        qr_path = salvar_qr_code(admin_user)
        
        # Mostrar informações
        print("\n=== Informações do Admin ===")
        print(f"Username: {admin_user.username}")
        print(f"Email: {admin_user.email}")
        print(f"Senha: {password}")
        print(f"Segredo OTP: {admin_user.otp_secret}")
        print(f"URI do OTP: {admin_user.get_otp_uri()}")
        if qr_path:
            print(f"QR Code salvo em: {qr_path}")
        else:
            print("QR Code não foi salvo. Use o URI do OTP ou o Segredo OTP para configuração manual.")
        
        print("\n=== Instruções para Configuração ===")
        print("1. Instale um aplicativo autenticador no seu celular")
        print("   (Google Authenticator, Microsoft Authenticator, etc)")
        print("2. Abra o aplicativo")
        print("3. Selecione a opção para adicionar uma nova conta")
        if qr_path:
            print("4. Escaneie o QR Code salvo em:", qr_path)
        print("\nOU configure manualmente:")
        print(f"- Nome da conta: {admin_user.username}")
        print(f"- Segredo OTP: {admin_user.otp_secret}")
        print("- Tipo: Baseado em tempo (TOTP)")
        print("- Algoritmo: SHA1")
        print("- Dígitos: 6")
        print("- Intervalo: 30 segundos")
        
        # Gerar código atual para verificação
        totp = pyotp.TOTP(admin_user.otp_secret)
        current_code = totp.now()
        print("\n=== Verificação do OTP ===")
        print(f"Código OTP atual: {current_code}")
        is_valid = admin_user.verify_otp(current_code)
        print(f"Verificação do código: {'Sucesso' if is_valid else 'Falha'}")

        if not is_valid:
            print("\nALERTA: Verificação do OTP falhou!")
            print("Por favor, verifique se o segredo OTP está correto.")
        
        # Fazer commit final para garantir que tudo foi salvo
        db.commit()
        
    except Exception as e:
        db.rollback()
        raise e
    finally:
        db.close()


def verify_admin(username=ADMIN_USERNAME, role=ADMIN_ROLE, save_qr=False):
    """Verifica se o usuário admin existe e tem OTP configurado"""
    db = get_db_session()
    try:
        admin_user = db.query(Usuario).filter_by(username=username).first()
        if admin_user is not None:
            print("\n=== Usuário Admin Encontrado ===")
            _ensure_admin_otp(db, admin_user)
            _ensure_admin_role(db, admin_user, role)
            return True
        else:
            print("\n=== Usuário Admin NÃO Encontrado ===")
            return False

    except Exception as e:
        print(f"Erro ao verificar o usuário admin: {e}")
        raise
    finally:
        db.close()


def rotate_admin_otp(username=ADMIN_USERNAME, save_qr=False):
    db = get_db_session()
    try:
        admin_user = db.query(Usuario).filter_by(username=username).first()
        if admin_user is None:
            print("Usuário admin não encontrado")
            return False

        admin_user.generate_otp_secret()
        db.commit()
        print(f"OTP do usuário '{username}' foi rotacionado.")
        print(f"Novo segredo OTP: {admin_user.otp_secret}")

        if save_qr:
            qr_path = salvar_qr_code(admin_user)
            if qr_path:
                print(f"Novo QR code salvo em: {qr_path}")
            else:
                print("Não foi possível salvar o QR code automaticamente.")

    except Exception:
        db.rollback()
        raise
    finally:
        db.close()


if __name__ == "__main__":
    create_admin()