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:
        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('admin_qr.png'),       # Diretório atual
        Path('/app/admin_qr.png')   # Diretório da aplicação
    ]
    
    # 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")
    
    # Tentar salvar em diferentes locais
    qr_saved = False
    saved_path = None
    
    for qr_path in qr_paths:
        try:
            # Tentar salvar o arquivo
            img.save(str(qr_path))
            print(f"QR Code salvo em: {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, 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 ===")
                if not admin.otp_secret:
                    print("Gerando novo segredo OTP...")
                    admin.generate_otp_secret()
                    db.commit()
            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()
                
                # Buscar ou criar role de admin
                admin_role = db.query(Role).filter_by(nome="admin").first()
                if not admin_role:
                    admin_role = Role(nome="admin", nivel=0)  # Nível 0 é o mais alto
                    db.add(admin_role)
                
                # Adicionar role ao usuário
                admin.roles.append(admin_role)
                
                # Adicionar e fazer commit
                db.add(admin)
                db.commit()
            
            # Gerar QR code
            qr_path, otp_uri = generate_qr_code(admin)
            
            if qr_path:
                print("\n=== QR Code Gerado ===")
                print(f"QR Code salvo em: {qr_path}")
                print(f"URI do OTP: {otp_uri}")
            else:
                print("\n=== QR Code Não Pode Ser Salvo ===")
                print("Use o URI OTP para configuração manual:")
                print(f"URI do OTP: {otp_uri}")
            
            # 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}")
            
            # 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")
            if qr_path:
                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()