Files
controles/create_admin.py

91 lines
3.2 KiB
Python

import os
from functions.database import get_db_connection, Usuario
from functions.rbac import Role
import pyotp
import qrcode
import base64
from io import BytesIO
def create_admin():
"""Cria o usuário admin se não existir"""
db = get_db_connection()
try:
# Verificar se o admin já existe
admin = db.query(Usuario).filter_by(username='admin').first()
if admin:
print("Usuário admin já existe")
# Verificar se o arquivo admin_qr.png existe
if os.path.exists('admin_qr.png'):
print("Usando OTP existente do arquivo admin_qr.png")
# Extrair o OTP secret do QR code existente
with open('admin_qr.png', 'rb') as f:
qr_data = f.read()
# Aqui você precisaria implementar a lógica para extrair o OTP secret do QR code
# Por enquanto, vamos apenas manter o OTP existente
return
else:
print("Gerando novo OTP para o admin...")
# Gerar novo OTP
otp_secret = pyotp.random_base32()
admin.otp_secret = otp_secret
db.commit()
else:
print("Criando usuário admin...")
# Criar usuário admin
admin = Usuario(
username='admin',
password='admin123',
is_admin=True
)
admin.email = 'admin@controles.com'
db.add(admin)
db.commit()
# Gerar OTP
otp_secret = pyotp.random_base32()
admin.otp_secret = otp_secret
db.commit()
# Atribuir role de Secretário Geral
admin_role = db.query(Role).filter_by(nivel=Role.SECRETARIO_GERAL).first()
if admin_role:
admin.roles.append(admin_role)
db.commit()
# Gerar QR code
totp = pyotp.TOTP(otp_secret)
provisioning_uri = totp.provisioning_uri(admin.username, issuer_name="Sistema de Controles")
qr = qrcode.QRCode(version=1, box_size=10, border=5)
qr.add_data(provisioning_uri)
qr.make(fit=True)
img = qr.make_image(fill_color="black", back_color="white")
# Salvar QR code como base64
buffered = BytesIO()
img.save(buffered, format="PNG")
qr_base64 = base64.b64encode(buffered.getvalue()).decode()
# Salvar QR code como arquivo
img.save('admin_qr.png')
print("\nConfiguração do OTP para o admin:")
print(f"OTP Secret: {otp_secret}")
print("\nInstruções:")
print("1. Use um aplicativo autenticador (como Google Authenticator ou Authy)")
print("2. Escaneie o QR code ou insira o OTP Secret manualmente")
print("3. Use o código gerado para fazer login")
print("\nQR code salvo em 'admin_qr.png'")
except Exception as e:
print(f"Erro ao criar admin: {str(e)}")
db.rollback()
raise
finally:
db.close()
if __name__ == '__main__':
create_admin()