bb-integrations-library 3.0.11__py3-none-any.whl

bb_integrations_lib/provider/ftp/interface.py

@@ -0,0 +1,273 @@
+import ftplib
+from abc import ABC
+from io import BytesIO, StringIO
+from pathlib import Path
+from stat import S_ISREG
+from typing import Iterable
+
+import paramiko
+import tenacity
+from loguru import logger
+from paramiko import RSAKey
+from tenacity import retry_if_exception_type, stop_after_attempt, wait_fixed
+
+from bb_integrations_lib.provider.ftp.model import FTPFileInfo, FTPAuthType, FTPType
+from bb_integrations_lib.secrets.credential_models import FTPCredential
+from bb_integrations_lib.shared.model import RawData, File
+
+
+class FTPClientInterface(ABC):
+    def __init__(self, credentials: FTPCredential):
+        self.credentials = credentials
+        self.first_connection = True
+
+    def connect(self):
+        self.first_connection = False
+
+    def disconnect(self):
+        raise NotImplementedError()
+
+    def _reconnect(self, retry_state: tenacity.RetryCallState):
+        raise NotImplementedError()
+
+    def cwd(self, directory: str) -> None:
+        raise NotImplementedError()
+
+    def list_files(self, directory: str) -> Iterable[str]:
+        raise NotImplementedError()
+
+    def rename_files(self, files: Iterable[tuple[str, str]]) -> None:
+        raise NotImplementedError()
+
+    def delete_files(self, paths: Iterable[str]) -> None:
+        raise NotImplementedError()
+
+    def upload_files(self, files: Iterable[File], path: str) -> None:
+        raise NotImplementedError()
+
+    def download_files(self, paths: Iterable[str]) -> Iterable[RawData]:
+        raise NotImplementedError()
+
+    @staticmethod
+    def reconnect_retry(func):
+        def wrap(*args, **kwargs):
+            if args[0].first_connection:
+                args[0].connect()
+            r = tenacity.Retrying(
+                reraise=True,
+                retry=retry_if_exception_type((OSError, AttributeError)),
+                stop=stop_after_attempt(3),
+                wait=wait_fixed(3),
+                after=args[0]._reconnect
+            )
+            for attempt in r:
+                with attempt:
+                    return func(*args, **kwargs)
+
+        return wrap
+
+class FTPClient(FTPClientInterface):
+    def __init__(self, credentials: FTPCredential):
+        super().__init__(credentials)
+
+        if self.credentials.ftp_type not in [FTPType.ftp, FTPType.ftps, FTPType.ftpes]:
+            raise NotImplementedError(
+                f"Attempted to use FTPClient with unsupported FTP type: {self.credentials.ftp_type} "
+                "(only supports ftp, ftps, ftpes)"
+            )
+
+        self.is_tls = self.credentials.ftp_type in [FTPType.ftps, FTPType.ftpes]
+        self.client: ftplib.FTP | None = None
+
+    def connect(self):
+        if self.is_tls:
+            self.client = ftplib.FTP_TLS(
+                host=self.credentials.host,
+                user=self.credentials.username,
+                passwd=self.credentials.password,
+            )
+            if self.credentials.ftp_type == FTPType.ftpes:
+                self.client.prot_p()
+        else:
+            self.client = ftplib.FTP(
+                host=self.credentials.host,
+                user=self.credentials.username,
+                passwd=self.credentials.password
+            )
+        super().connect()
+
+    def disconnect(self):
+        self.client.quit()
+        # Will force a reconnection next time it is used
+        self.first_connection = True
+
+    def _reconnect(self, retry_state: tenacity.RetryCallState):
+        # Targeted by reconnect_retry
+        self.connect()
+
+    @FTPClientInterface.reconnect_retry
+    def list_files(self, directory: str) -> Iterable[str]:
+        # Paths might be absolute - for consistency with the SFTP client, make them relative.
+        return [Path(x).name for x in self.client.nlst(directory)]
+
+    @FTPClientInterface.reconnect_retry
+    def rename_files(self, files: Iterable[tuple[str, str]]) -> None:
+        for old_name, new_name in files:
+            logger.debug(f"Renaming file {old_name} -> {new_name}")
+            self.client.rename(old_name, new_name)
+
+    @FTPClientInterface.reconnect_retry
+    def delete_files(self, paths: Iterable[str]) -> None:
+        for path in paths:
+            logger.debug(f"Deleting file {path}")
+            self.client.delete(path)
+
+    @FTPClientInterface.reconnect_retry
+    def upload_files(self, files: Iterable[File], path: str) -> None:
+        for file in files:
+            name = file.file_name
+            logger.debug(f"Uploading file {name}")
+            self.client.storbinary(f"STOR {path}/{name}", File.to_bytes(file.file_data))
+
+
+    @FTPClientInterface.reconnect_retry
+    def download_files(self, paths: Iterable[str]) -> Iterable[RawData]:
+        results = []
+        for path in paths:
+            logger.debug(f"Downloading {path}")
+            buf = BytesIO()
+            self.client.retrbinary(f"RETR {path}", buf.write)
+            buf.seek(0)
+            p = Path(path)
+            results.append(RawData(file_name=p.name, data=buf))
+        return results
+
+    @FTPClientInterface.reconnect_retry
+    def get_file_info(self, path: str) -> FTPFileInfo:
+        # Size is guaranteed
+        size = self.client.size(path)
+
+        return FTPFileInfo(
+            size=size,
+            permissions=None,
+            owner_id=None,
+            group_id=None,
+            last_access_time=None,
+            last_modification_time=None,
+        )
+
+
+class SFTPClient(FTPClientInterface):
+    def __init__(self, credentials: FTPCredential, private_key_path: str | None = None):
+        super().__init__(credentials)
+        self.ssh = paramiko.SSHClient()
+        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        self.sftp: paramiko.SFTPClient | None = None
+        self.private_key_path = private_key_path
+
+    def connect(self):
+        if self.credentials.private_key:
+            pkey = RSAKey.from_private_key(
+                StringIO(self.credentials.private_key, self.credentials.passphrase)
+            )
+        elif self.credentials.auth_type == FTPAuthType.rsa:
+            pkey = get_ppk()
+        else:
+            pkey = None
+
+
+        self.ssh.connect(
+            hostname=self.credentials.host,
+            username=self.credentials.username,
+            password=self.credentials.password,
+            port=self.credentials.port or 22,
+            key_filename=self.private_key_path,
+            pkey=pkey,
+            passphrase=self.credentials.passphrase,
+            look_for_keys=False,
+            allow_agent=False,
+        )
+        self.sftp = self.ssh.open_sftp()
+        super().connect()
+
+    def disconnect(self):
+        self.sftp.close()
+        self.first_connection = True
+
+    def _reconnect(self, retry_state: tenacity.RetryCallState):
+        self.connect()
+
+    @FTPClientInterface.reconnect_retry
+    def cwd(self, directory: str) -> None:
+        self.sftp.chdir(directory)
+
+    @FTPClientInterface.reconnect_retry
+    def list_files(self, directory: str) -> Iterable[str]:
+        logger.debug(f"Listing {directory}")
+        res = self.sftp.listdir_attr(directory)
+        return [x.filename for x in res if S_ISREG(x.st_mode)]
+
+
+    @FTPClientInterface.reconnect_retry
+    def rename_files(self, files: Iterable[tuple[str, str]]) -> None:
+        for old_path, new_path in files:
+            logger.debug(f"Renaming {old_path} -> {new_path}")
+            return self.sftp.rename(old_path, new_path)
+
+    @FTPClientInterface.reconnect_retry
+    def delete_files(self, paths: Iterable[str]) -> None:
+        for path in paths:
+            logger.debug(f"Deleting {path}")
+            self.sftp.remove(path)
+
+    @FTPClientInterface.reconnect_retry
+    def upload_files(self, files: Iterable[File], path: str) -> None:
+        for file in files:
+            name = file.file_name
+            logger.debug(f"Uploading file {name}")
+            self.sftp.putfo(file.to_bytes(file.file_data), f"{path}/{name}")
+
+    @FTPClientInterface.reconnect_retry
+    def download_files(self, paths: Iterable[str]) -> Iterable[RawData]:
+        results = []
+        for path in paths:
+            logger.debug(f"Downloading {path}")
+            buf = BytesIO()
+            self.sftp.getfo(path, buf)
+            buf.seek(0)
+            p = Path(path)
+            results.append(RawData(file_name=p.name, data=buf))
+        return results
+
+    @FTPClientInterface.reconnect_retry
+    def get_file_info(self, path: str) -> FTPFileInfo:
+        file_info = self.sftp.stat(path)
+        return FTPFileInfo(
+            size=file_info.st_size,
+            permissions=oct(file_info.st_mode) if file_info.st_mode else None,
+            owner_id=file_info.st_uid,
+            group_id=file_info.st_gid,
+            last_access_time=file_info.st_atime,
+            last_modification_time=file_info.st_mtime,
+        )
+
+def get_ppk(path: str = 'secrets/id_rsa') -> RSAKey:
+    """
+        Given a path, try to load RSA key file
+        :param path: The path to load
+        :return: The loaded RSA key
+    """
+    try:
+        private_key_path = Path(path)
+        if not private_key_path.exists():
+            raise FileNotFoundError(f"Private key file not found at: {private_key_path}")
+        private_key: RSAKey = RSAKey.from_private_key_file(str(private_key_path))
+        return private_key
+    except FileNotFoundError as fne:
+        msg = f'Failed to load private key from file: {path} -> {fne}'
+        logger.error(msg)
+        raise
+    except Exception as e:
+        msg = f'Failed to load private key from file: {path} -> {e}'
+        logger.error(msg)
+        raise

bb_integrations_lib/provider/ftp/model.py

@@ -0,0 +1,76 @@
+from datetime import datetime
+from enum import Enum
+from typing import Optional, Dict, Any
+from pydantic import BaseModel
+
+class FTPType(str, Enum):
+    """
+    Enumeration of supported FTP types.
+
+    Attributes:
+        ftps (str): FTPS (FTP over SSL) protocol.
+        sftp (str): SFTP (SSH File Transfer Protocol) protocol.
+        ftp (str): Standard FTP protocol.
+        ftpes (str): FTPeS (FTP over explicit SSL) protocol.
+    """
+    ftps = "ftps"
+    sftp = "sftp"
+    ftp = "ftp"
+    ftpes = "ftpes"
+
+
+class FTPAuthType(str, Enum):
+    """
+    Enumeration of FTP authentication methods.
+
+    Attributes:
+        basic (str): Basic authentication with username and password.
+        rsa (str): RSA authentication using private key.
+    """
+    basic = "basic"
+    rsa = "rsa"
+
+
+class FTPFileInfo(BaseModel):
+    """
+    Model representing metadata information about a file on the FTP server.
+
+    Attributes:
+        size (Optional[int]): Size of the file in bytes.
+        permissions (Optional[str]): File permissions.
+        owner_id (Optional[int]): User ID of the file owner.
+        group_id (Optional[int]): Group ID of the file owner.
+        last_access_time (Optional[int]): Timestamp of the last file access.
+        last_modification_time (Optional[int]): Timestamp of the last file modification.
+    """
+    size: Optional[int] = None
+    permissions: Optional[str] = None
+    owner_id: Optional[int] = None
+    group_id: Optional[int] = None
+    last_access_time: Optional[int] = None
+    last_modification_time: Optional[int] = None
+
+    @property
+    def last_modified_on(self):
+        """
+        Get the last modification date as a `datetime` object, if available.
+
+        Returns:
+            datetime or None: Last modified date if `last_modification_time` is set, otherwise None.
+        """
+        if self.last_modification_time:
+            return datetime.fromtimestamp(self.last_modification_time)
+        return None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert file info to a dictionary, including the `last_modified_on` property.
+
+        Returns:
+            Dict[str, Any]: Dictionary of file metadata with added `last_modified_on` field.
+        """
+        return {
+            **self.model_dump(),
+            "last_modified_on": self.last_modified_on,
+        }
+

bb_integrations_lib/provider/imap/client.py

@@ -0,0 +1,228 @@
+import email
+import imaplib
+import time
+from email.message import Message
+from io import BytesIO
+from typing import Iterable, Optional, List, Self
+
+import requests
+from charset_normalizer import detect
+from loguru import logger
+
+from bb_integrations_lib.secrets.credential_models import IMAPAuthOAuth, IMAPCredential
+from bb_integrations_lib.shared.model import RawData
+from bb_integrations_lib.util.utils import load_credentials
+
+
+class ImapEmailIntegration:
+    """Integration client to download attachments from an IMAP email server."""
+
+    def __init__(
+            self,
+            credentials: IMAPCredential,
+            criteria: str,
+            retries: int = 3,
+            type: str = 'sales',
+            file_extension: str = '.csv',
+    ):
+        self.credentials = credentials
+        self.criteria = criteria
+        self.retries = retries
+        self.type = type
+        self.file_extension = file_extension
+
+    def get_raw_data(self) -> Iterable[RawData]:
+        mail = IMAPClient(self.credentials)
+
+        mail.connect()
+        unseen_message_idxs = mail.search(self.criteria)
+        logger.info(f"📬 Found {len(unseen_message_idxs)} new emails in tank data inbox")
+
+        for idx in unseen_message_idxs:
+            retry_count = 0
+            while retry_count < self.retries:
+                try:
+                    message = mail.fetch(idx)
+                    attachment = mail.get_attachment_from_message(message, extension=self.file_extension)
+
+                    if not attachment:
+                        logger.error(f"⚠️ No valid attachment found in email id: {idx}")
+                        break
+
+                    if attachment is None:
+                        logger.error(f"⚠️ Could not process file from email id: {idx}")
+                        break  # No need to retry if processing failed
+
+                    raw_data = RawData(
+                        file_name=f"data_{idx}{self.file_extension}",
+                        data=attachment,
+                    )
+                    yield raw_data
+                    logger.info(f"✅ Successfully processed {self.file_extension} file from email id: {idx}")
+                    break  # Exit retry loop on success
+
+                except Exception as e:
+                    retry_count += 1
+                    logger.error(f"❌ Error reading data from email (Attempt {retry_count}/{self.retries}): {e}")
+                    time.sleep(5)
+                    if retry_count == self.retries:
+                        mail.mark_unseen(idx)
+
+
+class IMAPClient:
+    def __init__(
+            self,
+            credentials: IMAPCredential,
+            mailbox: str = "INBOX",
+            dry_run: bool = False,
+    ):
+        self.credentials = credentials
+        self.mailbox = mailbox
+        self.mail = None
+        self.dry_run = dry_run
+
+    @classmethod
+    def from_credential_file(cls, credential_file_name: str | None = None) -> Self:
+        credential_file_name = credential_file_name or "ftp.credentials"
+        credentials = load_credentials(credential_file_name)
+        return cls(credentials)
+
+    def test_imap(self, auth_string):
+        """Authenticate with Gmail IMAP using XOAUTH2 in dry_run mode."""
+        try:
+            imap_conn = imaplib.IMAP4_SSL('imap.gmail.com')
+            imap_conn.debug = 4
+            imap_conn.authenticate('XOAUTH2', lambda x: auth_string)
+            imap_conn.select('INBOX')
+            logger.success("✅ IMAP Authentication Successful! (Dry Run Mode)")
+            return True
+        except imaplib.IMAP4.error as e:
+            logger.error(f"❌ IMAP Authentication Failed (Dry Run Mode): {e}")
+            return False
+
+    def generate_oauth2_string(self, username, access_token):
+        """Generates a properly formatted OAuth2 authentication string."""
+        auth_string = f"user={username}\x01auth=Bearer {access_token}\x01\x01"
+        return auth_string.encode("utf-8")  # MUST return UTF-8 encoded bytes
+
+    def call_refresh_token(self, client_id, client_secret, refresh_token):
+        """Fetch a new access token using the refresh token."""
+        url = "https://oauth2.googleapis.com/token"
+        payload = {
+            "client_id": client_id,
+            "client_secret": client_secret,
+            "refresh_token": refresh_token,
+            "grant_type": "refresh_token"
+        }
+
+        headers = {"Content-Type": "application/x-www-form-urlencoded"}
+
+        response = requests.post(url, data=payload, headers=headers)
+
+        if response.status_code == 200:
+            return response.json()
+        else:
+            print("❌ Failed to refresh token:", response.text)
+            raise Exception(f"Failed to refresh token: {response.text}")
+
+    def refresh_authorization(self, client_id, client_secret, refresh_token):
+        """Refresh OAuth access token."""
+        response = self.call_refresh_token(client_id, client_secret, refresh_token)
+        if "access_token" in response:
+            return response["access_token"], response["expires_in"]
+        else:
+            raise Exception("❌ Failed to get access token: Invalid response")
+
+    def connect(self):
+        if self.mail:
+            return self.mail
+        self.mail = imaplib.IMAP4_SSL(self.credentials.host, self.credentials.port)
+        self.mail.debug = 4
+        auth = self.credentials.auth
+        if isinstance(auth, IMAPAuthOAuth):
+            try:
+                access_token, expires_in = self.refresh_authorization(
+                    client_id=auth.client_id,
+                    client_secret=auth.client_secret,
+                    refresh_token=auth.refresh_token
+                )
+                auth_string = self.generate_oauth2_string(self.credentials.email_address, access_token)
+                if self.dry_run:
+                    logger.info("Dry run mode: testing IMAP Authentication only.")
+                    self.test_imap(auth_string)
+                    return None
+                self.mail.authenticate("XOAUTH2", lambda x: auth_string)
+            except imaplib.IMAP4.error as e:
+                logger.error(f"Gmail authentication failed: {e}")
+                raise Exception("Invalid Gmail OAuth authentication. Check your credentials and refresh token.")
+        else:
+            self.mail.login(self.credentials.email_address, self.credentials.auth.password)
+        self.mail.select(self.mailbox)
+        return self.mail
+
+    def change_mailbox(self, mailbox: str):
+        """Change the mailbox (folder) in IMAP."""
+        self.connect().select(mailbox)
+
+    def search(self, criteria: str) -> List[int]:
+        """Search for emails based on criteria (e.g., '(UNSEEN)')."""
+        status, message_indexes = self.connect().search(None, criteria)
+        message_indexes = [int(idx) for idx in next(iter(message_indexes), "").split()]
+        return message_indexes
+
+    def fetch(self, message_index: int) -> Optional[Message]:
+        """Fetch an email by its index."""
+        status, message = self.connect().fetch(str(message_index), "(RFC822)")
+        try:
+            if message and message[0]:
+                raw_email_string = message[0][1].decode("utf-8")
+                return email.message_from_string(raw_email_string)
+        except (IndexError, AttributeError) as e:
+            logger.error(f"Error fetching email {message_index}: {e}")
+        return None
+
+    def mark_unseen(self, message_index: int):
+        """Mark an email as unseen (unread)."""
+        try:
+            self.connect().store(str(message_index), "-FLAGS", "\\Seen")
+        except Exception as e:
+            logger.error(f"Error marking email {message_index} as unseen: {e}")
+
+    @staticmethod
+    def get_attachment_from_message(message: Message, extension: str, return_rawdata: bool = False) -> Optional[bytes | str | RawData]:
+        """Extracts an email attachment and handles unknown encoding.
+
+        - Returns **decoded text** for text-based files (CSV, TXT, JSON, etc.).
+        - Returns **raw binary data** for non-text files (PDF, XLSX, ZIP, etc.).
+        """
+        for part in message.walk():
+            if part.get_content_maintype() == 'multipart':
+                continue
+            if part.get("Content-Disposition") is None and part.get_filename() is None:
+                continue
+            if part.get_filename() and part.get_filename().lower().endswith(extension.lower()):
+                try:
+                    raw_bytes = part.get_payload(decode=True)
+                    if extension.lower() in [".pdf", ".xlsx", ".xls", ".zip", ".png", ".jpg", ".jpeg", ".gif", ".docx",
+                                             ".pptx"]:
+                        logger.info(f"🗂 Binary file detected ({part.get_filename()}), returning raw bytes")
+                        return raw_bytes
+                    detected = detect(raw_bytes)
+                    encoding = detected["encoding"]
+                    if not encoding:
+                        logger.warning(
+                            f"⚠️ Encoding could not be detected for {part.get_filename()}, defaulting to latin-1")
+                        encoding = "latin-1"
+                    logger.info(f"📄 Detected encoding: {encoding} for {part.get_filename()}")
+                    if return_rawdata:
+                        return RawData(file_name=part.get_filename(), data=BytesIO(raw_bytes))
+                    else:
+                        return raw_bytes.decode(encoding, errors="replace")  # ✅ Decode with replacement for errors
+                except UnicodeDecodeError as e:
+                    logger.error(f"❌ Failed to decode {part.get_filename()} as text: {e}")
+                    if return_rawdata:
+                        return RawData(file_name=part.get_filename(), data=BytesIO(raw_bytes))
+                    else:
+                        return raw_bytes  # ✅ Return raw bytes if decoding fails
+        return None
+

bb_integrations_lib/provider/imap/model.py

@@ -0,0 +1,3 @@
+
+
+

bb_integrations_lib/provider/sqlserver/client.py

@@ -0,0 +1,106 @@
+from typing import Dict, List, Any, Optional
+from contextlib import contextmanager
+from urllib.parse import quote_plus
+
+from loguru import logger
+from sqlalchemy import create_engine, text, MetaData
+from sqlalchemy.engine import Engine, Connection, Result
+from sqlalchemy.orm import sessionmaker, Session
+
+
+class SQLServerClient:
+    def __init__(
+            self,
+            server: str,
+            database: str,
+            username: Optional[str] = None,
+            password: Optional[str] = None,
+            driver: str = "ODBC Driver 17 for SQL Server",
+            trusted_connection: bool = False,
+            echo: bool = False,
+            mars_connection: bool = True
+    ):
+
+        self.server = server
+        self.database = database
+        self.username = username
+        self.password = password
+        self.driver = driver
+        self.trusted_connection = trusted_connection
+        self.echo = echo
+        self.mars_connection = mars_connection
+        self.engine = self._create_engine()
+        self.metadata = MetaData()
+        self.Session = sessionmaker(bind=self.engine)
+
+    def _create_engine(self) -> Engine:
+        connection_string = self._build_connection_string()
+        return create_engine(connection_string, echo=self.echo)
+
+    def _build_connection_string(self) -> str:
+        params = {
+            "driver": self.driver,
+            "server": self.server,
+            "database": self.database,
+        }
+
+        if self.trusted_connection:
+            params["trusted_connection"] = "yes"
+        else:
+            params["uid"] = self.username
+            params["pwd"] = self.password
+        if self.mars_connection:
+            params["Mars_Connection"] = "yes"
+        params["TrustServerCertificate"] = "yes"
+        params["MultipleActiveResultSets"] = "True" if self.mars_connection else "False"
+        params["Connection Timeout"] = "30"
+        params["Command Timeout"] = "30"
+        connection_string_parts = []
+        for key, value in params.items():
+            if value is not None:
+                connection_string_parts.append(f"{key}={quote_plus(str(value))}")
+
+        connection_string = ";".join(connection_string_parts)
+        return f"mssql+pyodbc:///?odbc_connect={quote_plus(connection_string)}"
+
+    @contextmanager
+    def get_connection(self) -> Connection:
+        connection = self.engine.connect()
+        try:
+            yield connection
+        finally:
+            connection.close()
+
+    @contextmanager
+    def get_session(self) -> Session:
+        session = self.Session()
+        try:
+            yield session
+            session.commit()
+        except:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+    def execute_query(self, query: str, params: Optional[Dict[str, Any]] = None) -> Result:
+        with self.get_session() as session:
+            result = session.execute(text(query), params or {})
+            return result
+
+    def get_all(self, query: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
+        """
+        :rtype: object
+        """
+        with self.get_session() as session:
+            result = session.execute(text(query), params or {})
+            rows = result.fetchall()
+            return [dict(row._mapping) for row in rows]
+
+    def get_mappings(self, query: str, params: Optional[Dict[str, Any]] = None, source_system: Optional[str] = None,
+                     mapping_type: Optional[str] = None) -> List[Dict[str, Any]]:
+        if source_system is not None:
+            logger.warning("Source System not implemented")
+        if mapping_type is not None:
+            logger.warning("Mapping Type not implemented")
+        return self.get_all(query, params or {})