From 5f10132f0dabbecbf66e1deede80ec331d147d24 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Magalh=C3=A3es?= <joamag@gmail.com>
Date: Tue, 1 Oct 2024 14:38:16 +0100
Subject: [PATCH] fix: issue with investment extracts

---
 CHANGELOG.md               |  2 +-
 src/weby_pilot/__main__.py |  4 ++--
 src/weby_pilot/base.py     |  7 +++++++
 src/weby_pilot/bpi.py      | 11 +++++------
 4 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 373860c..77f3e7c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,7 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
-*
+* Issue with investment extracts
 
 ## [0.1.4] - 2024-10-01
 
diff --git a/src/weby_pilot/__main__.py b/src/weby_pilot/__main__.py
index eaee058..d819f49 100644
--- a/src/weby_pilot/__main__.py
+++ b/src/weby_pilot/__main__.py
@@ -3,8 +3,8 @@
 
 from .bpi import BpiAPI
 
-print(BpiAPI().download_account_report(report_indexes=range(0, 2), account_index=1))
-
+print(BpiAPI().download_account_report(report_indexes=range(0, 2)))
+# print(BpiAPI().download_account_report(report_indexes=range(0, 2), account_index=1))
 # BpiAPI().download_report(section="Extrato Investimento", report_indexes=range(0, 8))
 # BpiAPI().download_invoice()
 # print(BpiAPI().get_balance())
diff --git a/src/weby_pilot/base.py b/src/weby_pilot/base.py
index 5f4bbf4..3990c25 100644
--- a/src/weby_pilot/base.py
+++ b/src/weby_pilot/base.py
@@ -2,6 +2,7 @@
 # -*- coding: utf-8 -*-
 
 from io import BytesIO
+from shutil import rmtree
 from typing import IO, Any, Generator, List, Literal
 from uuid import uuid4
 from time import sleep
@@ -30,6 +31,7 @@ def __init__(self, headless: bool = False, stable: bool = True):
 class WebyAPI:
     _driver: Chrome | None = None
     _wait: WebDriverWait[Chrome] | None = None
+    _remove_downloads: bool = False
     _downloads_dir: str = abspath("downloads")
     _temp_dir: str = abspath("temp")
     _last_path: str | None = None
@@ -64,13 +66,18 @@ def start(self, options: WebyOptions = WebyOptions()):
                 "safebrowsing.disable_download_protection": True,
             },
         )
+
         self._driver = Chrome(options=chrome_options)
+        if options.stable:
+            self._driver.delete_all_cookies()
 
         self._wait = WebDriverWait(self._driver, MAX_WAIT_TIME)
 
     def stop(self):
         if self.driver is not None:
             self.driver.quit()
+        if self._remove_downloads and exists(self._downloads_dir):
+            rmtree(self._downloads_dir)
 
     def get_element(self, by: str, value: str) -> WebElement:
         return self.wait.until(
diff --git a/src/weby_pilot/bpi.py b/src/weby_pilot/bpi.py
index 684aba2..3f5f487 100644
--- a/src/weby_pilot/bpi.py
+++ b/src/weby_pilot/bpi.py
@@ -121,14 +121,15 @@ def download_report(
         self,
         section: ReportSections = "Extrato Conta",
         report_indexes: Sequence[int] = (0,),
-        account_index: int = 0,
+        account_index: int | None = None,
     ) -> Sequence["BpiDocument"]:
         docs: list[BpiDocument] = []
         with self.driver_ctx():
             self.login()
             self.select_section("Consultas")
             self.select_side_menu(cast(BpiSideSections, section))
-            self.select_account(account_index)
+            if account_index is not None:
+                self.select_account(account_index)
             for report_index in report_indexes:
                 self.click_extract(row_index=report_index)
                 docs.append(
@@ -155,12 +156,10 @@ def download_account_report(
         )
 
     def download_investing_report(
-        self, report_indexes: Sequence[int] = (0,), account_index: int = 0
+        self, report_indexes: Sequence[int] = (0,)
     ) -> Sequence["BpiDocument"]:
         return self.download_report(
-            section="Extrato Investimento",
-            report_indexes=report_indexes,
-            account_index=account_index,
+            section="Extrato Investimento", report_indexes=report_indexes
         )
 
     def download_card_report(