155 lines
4.6 KiB
Python
155 lines
4.6 KiB
Python
# smb_excel.py (обновлённый фрагмент)
|
||
from io import BytesIO
|
||
from typing import Optional, Union, Any, Dict
|
||
#from smb.SMBConnection import SMBConnection
|
||
import pandas as pd
|
||
import re
|
||
|
||
class SMBDownloadError(Exception):
|
||
pass
|
||
|
||
def _normalize_path(p: str) -> str:
|
||
"""Убирает случайные r"..."/"...", лишние кавычки и ставит прямые слэши."""
|
||
if p is None:
|
||
return p
|
||
# уберём обрамляющие кавычки/префикс r"..."
|
||
m = re.fullmatch(r"""r?["'](.*)["']""", p.strip())
|
||
if m:
|
||
p = m.group(1)
|
||
# заменим \ на /
|
||
p = p.replace("\\", "/").lstrip("/") # путь внутри шары не должен начинаться с /
|
||
return p
|
||
|
||
def _try_connect(
|
||
remote_name: str,
|
||
server_ip: str,
|
||
username: str,
|
||
password: str,
|
||
*,
|
||
client_machine_name: str,
|
||
domain: str,
|
||
port: int
|
||
) -> Optional[SMBConnection]:
|
||
conn = SMBConnection(
|
||
username,
|
||
password,
|
||
client_machine_name,
|
||
remote_name, # ВАЖНО: remote_name — имя сервера (или IP как fallback)
|
||
domain=domain,
|
||
use_ntlm_v2=True,
|
||
is_direct_tcp=True
|
||
)
|
||
try:
|
||
if conn.connect(server_ip, port):
|
||
return conn
|
||
except Exception:
|
||
pass
|
||
return None
|
||
|
||
def _connect_smb(
|
||
server_name: str,
|
||
server_ip: str,
|
||
username: str,
|
||
password: str,
|
||
*,
|
||
client_machine_name: str = "client",
|
||
domain: str = "",
|
||
port: int = 445,
|
||
) -> SMBConnection:
|
||
"""
|
||
Пробуем два варианта remote_name: (1) реальное имя сервера, (2) IP.
|
||
"""
|
||
for remote_name in (server_name, server_ip):
|
||
if not remote_name:
|
||
continue
|
||
conn = _try_connect(
|
||
remote_name=remote_name,
|
||
server_ip=server_ip,
|
||
username=username,
|
||
password=password,
|
||
client_machine_name=client_machine_name,
|
||
domain=domain,
|
||
port=port,
|
||
)
|
||
if conn:
|
||
return conn
|
||
raise SMBDownloadError(
|
||
f"Не удалось подключиться к SMB {server_ip}:{port}. "
|
||
f"Проверь server_name (реальное имя хоста), домен/логин и доступность порта 445."
|
||
)
|
||
|
||
def fetch_smb_file(
|
||
*,
|
||
server_name: str,
|
||
server_ip: str,
|
||
share_name: str,
|
||
file_path_in_share: str,
|
||
username: str,
|
||
password: str,
|
||
client_machine_name: str = "client",
|
||
domain: str = "",
|
||
port: int = 445,
|
||
) -> BytesIO:
|
||
"""
|
||
Возвращает содержимое файла из SMB как BytesIO.
|
||
"""
|
||
file_path_in_share = _normalize_path(file_path_in_share)
|
||
conn: Optional[SMBConnection] = None
|
||
try:
|
||
conn = _connect_smb(
|
||
server_name=server_name,
|
||
server_ip=server_ip,
|
||
username=username,
|
||
password=password,
|
||
client_machine_name=client_machine_name,
|
||
domain=domain,
|
||
port=port,
|
||
)
|
||
# Быстрая проверка существования каталога/файла
|
||
parent = "/".join(file_path_in_share.split("/")[:-1]) or ""
|
||
# listPath кидает исключение, если каталога/шары нет — получим понятную ошибку
|
||
if parent:
|
||
conn.listPath(share_name, parent)
|
||
|
||
buf = BytesIO()
|
||
conn.retrieveFile(share_name, file_path_in_share, buf)
|
||
buf.seek(0)
|
||
return buf
|
||
except Exception as e:
|
||
raise SMBDownloadError(
|
||
f"Ошибка скачивания {share_name}/{file_path_in_share}: {e}"
|
||
) from e
|
||
finally:
|
||
if conn:
|
||
try:
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
|
||
def read_excel_from_smb(
|
||
*,
|
||
server_name: str,
|
||
server_ip: str,
|
||
share_name: str,
|
||
file_path_in_share: str,
|
||
username: str,
|
||
password: str,
|
||
client_machine_name: str = "client",
|
||
domain: str = "",
|
||
port: int = 445,
|
||
sheet_name: Optional[Union[str, int, list]] = None,
|
||
**read_excel_kwargs: Dict[str, Any],
|
||
) -> pd.DataFrame:
|
||
file_obj = fetch_smb_file(
|
||
server_name=server_name,
|
||
server_ip=server_ip,
|
||
share_name=share_name,
|
||
file_path_in_share=file_path_in_share,
|
||
username=username,
|
||
password=password,
|
||
client_machine_name=client_machine_name,
|
||
domain=domain,
|
||
port=port,
|
||
)
|
||
return pd.read_excel(file_obj, sheet_name=sheet_name, **read_excel_kwargs)
|