controles/create_admin.py

from functions.database import init_database, Usuario, Role, get_db_connection
import qrcode
import os
from pathlib import Path
import pyotp

def generate_qr_code(user):
    """
    Gera o QR code para um usuário específico
    
    Args:
        user: Instância do modelo Usuario
        
    Returns:
        Path: Caminho do arquivo QR code gerado
    """
    # Gerar QR Code apenas na raiz do projeto
    qr_path = Path('admin_qr.png')
    
    # Remover arquivo antigo se existir
    if qr_path.exists():
        os.remove(str(qr_path))
    
    # Gerar e salvar QR Code
    qr = qrcode.QRCode(version=1, box_size=10, border=5)
    
    # Gerar URI do OTP
    totp = pyotp.TOTP(user.otp_secret)
    otp_uri = totp.provisioning_uri(
        name=user.username,
        issuer_name="Sistema de Controles"
    )
    
    qr.add_data(otp_uri)
    qr.make(fit=True)
    img = qr.make_image(fill_color="black", back_color="white")
    img.save(str(qr_path))
    
    print(f"\nQR Code gerado em: {os.path.abspath(qr_path)}")
    
    return qr_path, otp_uri

def create_admin_user():
    """Cria ou atualiza o usuário admin"""
    try:
        # Inicializar banco de dados
        init_database()
        
        # Criar sessão
        db = get_db_connection()
        
        try:
            # Verificar se já existe um usuário admin
            admin = db.query(Usuario).filter_by(username="admin").first()
            
            if admin:
                print("\n=== Usuário Admin Encontrado ===")
            else:
                print("\n=== Criando Novo Usuário Admin ===")
                # Criar novo usuário admin
                admin = Usuario(
                    username="admin",
                    email="admin@example.com",
                    is_admin=True
                )
                admin.set_password("admin123")
                admin.generate_otp_secret()
                
                # Adicionar e fazer commit
                db.add(admin)
                db.commit()
            
            # Gerar QR code uma única vez
            qr_path, otp_uri = generate_qr_code(admin)
            
            # Mostrar informações
            print("\n=== Informações do Admin ===")
            print(f"Username: {admin.username}")
            print(f"Email: {admin.email}")
            print(f"Senha: admin123")
            print(f"Segredo OTP: {admin.otp_secret}")
            
            print("\n=== QR Code Gerado ===")
            print(f"QR Code salvo em: {qr_path}")
            print(f"URI do OTP: {otp_uri}")
            
            # Gerar código atual para verificação
            totp = pyotp.TOTP(admin.otp_secret)
            current_code = totp.now()
            print("\n=== Verificação do OTP ===")
            print(f"Código OTP atual: {current_code}")
            print(f"Verificação do código: {totp.verify(current_code)}")
            
            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")
            print("4. Escaneie o QR Code salvo em:", qr_path)
            print("\nOU configure manualmente:")
            print(f"- Nome da conta: {admin.username}")
            print(f"- Segredo: {admin.otp_secret}")
            print("- Tipo: Baseado em tempo (TOTP)")
            print("- Algoritmo: SHA1")
            print("- Dígitos: 6")
            print("- Intervalo: 30 segundos")
            
            # Verificação final
            print("\n=== Teste de Verificação ===")
            test_code = totp.now()
            print(f"Código de teste: {test_code}")
            is_valid = admin.verify_otp(test_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()
            
    except Exception as e:
        print(f"\nErro durante a execução: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    create_admin_user()