controles/create_test_users.py

from functions.database import get_db_connection, Usuario
from functions.rbac import Role
import pyotp
import qrcode
import os
import base64
from io import BytesIO

def create_test_users():
    """Cria usuários de teste se não existirem"""
    db = get_db_connection()
    try:
        # Usuários de teste
        test_users = [
            {
                'username': 'teste',
                'password': 'admin123',  # Mesma senha do admin
                'email': 'teste@controles.com',
                'is_admin': True
            },
            {
                'username': 'aligner',
                'password': 'Test123!@#',
                'email': 'aligner@controles.com',
                'is_admin': False
            },
            {
                'username': 'tester',
                'password': 'Test123!@#',
                'email': 'tester@controles.com',
                'is_admin': False
            },
            {
                'username': 'deployer',
                'password': 'Test123!@#',
                'email': 'deployer@controles.com',
                'is_admin': False
            }
        ]
        
        # Obter o OTP secret do admin se existir
        admin = db.query(Usuario).filter_by(username='admin').first()
        admin_otp_secret = admin.otp_secret if admin else None
        
        for user_data in test_users:
            # Verificar se o usuário já existe
            user = db.query(Usuario).filter_by(username=user_data['username']).first()
            
            if not user:
                print(f"Criando usuário {user_data['username']}...")
                # Criar usuário
                user = Usuario(
                    username=user_data['username'],
                    email=user_data['email'],
                    is_admin=user_data['is_admin']
                )
                user.set_password(user_data['password'])
                user.tipo = "ADMIN" if user_data['is_admin'] else "USUARIO"
                
                # Se for o usuário teste, usar o mesmo OTP do admin
                if user_data['username'] == 'teste' and admin_otp_secret:
                    user.otp_secret = admin_otp_secret
                else:
                    # Gerar novo OTP para outros usuários
                    user.otp_secret = pyotp.random_base32()
                
                db.add(user)
                db.commit()
                
                # Atribuir role de Secretário Geral para o usuário teste
                if user_data['username'] == 'teste':
                    admin_role = db.query(Role).filter_by(nivel=Role.SECRETARIO_GERAL).first()
                    if admin_role:
                        user.roles.append(admin_role)
                        db.commit()
                
                print(f"Usuário {user_data['username']} criado com sucesso!")
                
                # Gerar QR code para o novo usuário
                qr_path = f"{user_data['username']}_qr.png"
                if not os.path.exists(qr_path):
                    totp = pyotp.TOTP(user.otp_secret)
                    qr = qrcode.QRCode(version=1, box_size=10, border=5)
                    qr.add_data(totp.provisioning_uri(user.email, issuer_name="Sistema de Controles"))
                    qr.make(fit=True)
                    img = qr.make_image(fill_color="black", back_color="white")
                    img.save(qr_path)
                    print(f"QR Code gerado para {user_data['username']} em: {qr_path}")
            else:
                print(f"Usuário {user_data['username']} já existe")
                
                # Se for o usuário teste e não tiver o OTP do admin, atualizar
                if user_data['username'] == 'teste' and admin_otp_secret and user.otp_secret != admin_otp_secret:
                    user.otp_secret = admin_otp_secret
                    db.commit()
                    print(f"OTP do usuário teste atualizado para o mesmo do admin")
                elif not user.otp_secret:
                    # Se não tiver OTP, gerar um novo
                    user.otp_secret = pyotp.random_base32()
                    db.commit()
                    print(f"Novo OTP gerado para {user_data['username']}")
                    
                    # Gerar QR code
                    qr_path = f"{user_data['username']}_qr.png"
                    if not os.path.exists(qr_path):
                        totp = pyotp.TOTP(user.otp_secret)
                        qr = qrcode.QRCode(version=1, box_size=10, border=5)
                        qr.add_data(totp.provisioning_uri(user.email, issuer_name="Sistema de Controles"))
                        qr.make(fit=True)
                        img = qr.make_image(fill_color="black", back_color="white")
                        img.save(qr_path)
                        print(f"QR Code gerado para {user_data['username']} em: {qr_path}")
                
                # Verificar se o usuário teste tem a role de Secretário Geral
                if user_data['username'] == 'teste':
                    admin_role = db.query(Role).filter_by(nivel=Role.SECRETARIO_GERAL).first()
                    if admin_role and admin_role not in user.roles:
                        user.roles.append(admin_role)
                        db.commit()
                        print(f"Role de Secretário Geral atribuída ao usuário teste")
    
    except Exception as e:
        print(f"Erro ao criar usuários de teste: {str(e)}")
        db.rollback()
        raise
    finally:
        db.close()

if __name__ == '__main__':
    create_test_users()