import os import sys import jwt import psutil import requests import psycopg2 from psycopg2.extras import DictCursor import logging from pathlib import Path from http.cookies import SimpleCookie from cryptography.fernet import Fernet from datetime import datetime, timedelta from pytz import timezone # --- Configuração de Logging --- LOG_DIR = Path("/app/logs") LOG_DIR.mkdir(exist_ok=True) logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[ logging.FileHandler(LOG_DIR / "worker.log"), logging.StreamHandler(sys.stdout) ] ) # --- Configurações Iniciais --- APP_KEY = os.getenv("APP_KEY") DB_NAME = os.getenv("POSTGRES_DB") DB_USER = os.getenv("POSTGRES_USER") DB_PASS = os.getenv("POSTGRES_PASSWORD") DB_HOST = os.getenv("POSTGRES_HOST") SAO_PAULO_TZ = timezone('America/Sao_Paulo') RENEW_URL = "https://fgtsdigital.sistema.gov.br/portal/mensagem/v1/mensagens/somamsgsnaolidas" MAX_RETRIES = 3 LOCK_FILE = Path("/tmp/worker.lock") # --- Módulo de Criptografia --- if not APP_KEY: logging.error("APP_KEY não definida no ambiente.1") raise ValueError("APP_KEY não definida no ambiente.") fernet = Fernet(APP_KEY.encode()) def decrypt_data(encrypted_data: str) -> str: """Descriptografa dados usando Fernet.""" return fernet.decrypt(encrypted_data.encode()).decode() def encrypt_data(data: str) -> str: """Criptografa dados usando Fernet.""" return fernet.encrypt(data.encode()).decode() # --- Funções Auxiliares --- def get_db_connection(): """Estabelece conexão com o PostgreSQL.""" try: return psycopg2.connect(dbname=DB_NAME, user=DB_USER, password=DB_PASS, host=DB_HOST, port="5432") except psycopg2.OperationalError as e: logging.error(f"Falha ao conectar ao banco de dados: {e}") return None def process_session(session_id: int, cnpj: str, encrypted_cookie: str): """Lida com a renovação de uma única sessão usando uma fusão de dicionários para evitar duplicatas.""" logging.info(f"Iniciando processo de renovação para CNPJ {cnpj} (ID: {session_id}).") conn = get_db_connection() if not conn: return for attempt in range(MAX_RETRIES): try: # 1. Descriptografar o cookie completo do banco de dados original_cookie_str = decrypt_data(encrypted_cookie) # 2. Converter a string de cookie antiga em um dicionário. # Esta é a nossa fonte da verdade para os cookies que devem ser preservados. old_cookies_morsels = SimpleCookie() old_cookies_morsels.load(original_cookie_str) old_cookies_dict = {key: morsel.value for key, morsel in old_cookies_morsels.items()} # 3. Fazer a requisição de renovação, passando os cookies antigos. # Não precisamos mais de um objeto Session complexo. headers = { 'Accept': 'application/json, text/plain, */*', 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36', } logging.info(f"[{cnpj}] Tentativa {attempt + 1}: Fazendo requisição para o endpoint de renovação...") response = requests.get(RENEW_URL, headers=headers, cookies=old_cookies_dict, timeout=60) response.raise_for_status() # 4. Obter os cookies NOVOS da resposta como um dicionário. new_cookies_dict = response.cookies.get_dict() # 5. ATUALIZAR o dicionário antigo com os valores do novo. # Esta é a etapa crucial. dict.update() sobrescreve as chaves existentes. merged_cookies_dict = old_cookies_dict.copy() merged_cookies_dict.update(new_cookies_dict) # 6. Reconstruir a string de cookie final a partir do dicionário mesclado. new_full_cookie_str = "; ".join([f"{key}={value}" for key, value in merged_cookies_dict.items()]) logging.debug(f"[{cnpj}] Cookie mesclado com sucesso. String final: {new_full_cookie_str}") # 7. Extrair novo token e data de expiração do DICIONÁRIO mesclado. new_fgtsd_token_str = merged_cookies_dict.get('fgtsd_token') if not new_fgtsd_token_str: raise ValueError("O 'fgtsd_token' não foi encontrado no dicionário de cookies mesclado.") decoded_token = jwt.decode(new_fgtsd_token_str, options={"verify_signature": False}) new_exp_timestamp = decoded_token.get("exp") new_exp_datetime = datetime.fromtimestamp(new_exp_timestamp, tz=SAO_PAULO_TZ) # 8. Criptografar e atualizar o banco de dados new_encrypted_cookie = encrypt_data(new_full_cookie_str) with conn.cursor() as cur: query = """ UPDATE companies SET cookie = %s, expires_on = %s, updated_at = %s, status = 'active', details = NULL WHERE id = %s; """ cur.execute(query, (new_encrypted_cookie, new_exp_datetime, datetime.now(SAO_PAULO_TZ), session_id)) conn.commit() logging.info(f"SUCESSO: Sessão para o CNPJ {cnpj} renovada. Nova expiração: {new_exp_datetime.strftime('%Y-%m-%d %H:%M:%S')}") return except Exception as e: # O erro que você viu anteriormente foi capturado aqui logging.error(f"FALHA (Tentativa {attempt + 1}/{MAX_RETRIES}) para CNPJ {cnpj}: {e}") if attempt + 1 == MAX_RETRIES: error_message = f"Falha na renovação após {MAX_RETRIES} tentativas. Último erro: {e}" with conn.cursor() as cur: cur.execute("UPDATE companies SET status = 'with_error', details = %s WHERE id = %s;", (error_message, session_id)) conn.commit() finally: if conn: conn.close() def is_process_running(pid: int) -> bool: """Verifica se o processo com o PID especificado está rodando.""" try: return psutil.Process(pid).is_running() except (psutil.NoSuchProcess, psutil.AccessDenied): return False def main(): """Função principal que busca e processa as sessões.""" logging.info("Worker iniciado. Verificando por sessões a serem renovadas.") conn = get_db_connection() if not conn: return try: with conn.cursor(cursor_factory=DictCursor) as cur: # Seleciona sessões ativas que expiram nos próximos 10 minutos # A comparação de tempo é feita em UTC para evitar problemas de fuso ten_minutes_from_now_utc = datetime.now(timezone('UTC')) + timedelta(minutes=10) query = """ SELECT id, cnpj, cookie FROM companies WHERE status = 'active' AND expires_on <= %s ORDER BY expires_on ASC LIMIT 10; """ cur.execute(query, (ten_minutes_from_now_utc,)) sessions_to_renew = cur.fetchall() if not sessions_to_renew: logging.info("Nenhuma sessão para renovar no momento.") return logging.info(f"Encontradas {len(sessions_to_renew)} sessões para renovar.") for session in sessions_to_renew: process_session(session['id'], session['cnpj'], session['cookie']) except Exception as e: logging.error(f"Erro geral no loop principal do worker: {e}") finally: conn.close() # Remove o arquivo de lock ao final da execução LOCK_FILE.unlink() logging.info("Worker finalizado. Lock liberado.") if __name__ == "__main__": try: # Verifica se existe um lock anterior if LOCK_FILE.exists(): try: with LOCK_FILE.open("r") as f: old_pid = int(f.read().strip()) # Se o processo não está mais rodando, remove o lock órfão if not is_process_running(old_pid): logging.warning(f"Lock órfão detectado (PID {old_pid} não está rodando). Removendo...") LOCK_FILE.unlink() else: logging.warning(f"Execução do worker já em andamento (PID {old_pid}). Saindo.") sys.exit(0) except (ValueError, IOError) as e: logging.warning(f"Erro ao ler lock file: {e}. Removendo lock corrompido...") LOCK_FILE.unlink() # Cria novo lock with LOCK_FILE.open("x") as f: f.write(str(os.getpid())) main() except FileExistsError: logging.warning(f"Lock já existe em {LOCK_FILE}. Saindo.") sys.exit(0) except Exception as e: logging.critical(f"Erro crítico não capturado no worker: {e}") if LOCK_FILE.exists(): LOCK_FILE.unlink()