Skip to content

Commit 97e596f

Browse files
spuiukspuiuk
committed
testhelper/smbclient: use smbprotocol python module
pysmb only support SMB2.0 version. The samba-container defaults to using SMB2.1 protocol and we are unable to use pysmb to connect to the samba server. This is a major issue. The smbprotocol python module avoids this problem. Signed-off-by: Sachin Prabhu <[email protected]>
1 parent a11493f commit 97e596f

File tree

2 files changed

+97
-71
lines changed

2 files changed

+97
-71
lines changed

‎testhelper/smbclient.py

Lines changed: 94 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -1,111 +1,137 @@
1-
from smb.SMBConnection import SMBConnection # type: ignore
2-
from smb import smb_structs, base # type: ignore
1+
from smbprotocol.exceptions import SMBException # type: ignore
2+
import smbclient # type: ignore
33
import typing
4-
import io
4+
5+
rw_chunk_size = 1 << 21 # 2MB
56

67

78
class SMBClient:
8-
"""Use pysmb to access the SMB server"""
9+
"""Use smbprotocol python module to access the SMB server"""
910

10-
def __init__(self, hostname: str, share: str, username: str, passwd: str):
11+
def __init__(
12+
self,
13+
hostname: str,
14+
share: str,
15+
username: str,
16+
passwd: str,
17+
port: int = 445,
18+
):
1119
self.server = hostname
1220
self.share = share
13-
self.username = username
14-
self.password = passwd
21+
self.port = port
22+
self.client_params = {
23+
"username": username,
24+
"password": passwd,
25+
"connection_cache": {},
26+
}
27+
self.prepath = f"\\\\{self.server}\\{self.share}\\"
1528
self.connected = False
1629
self.connect()
1730

31+
def _path(self, path: str = "/") -> str:
32+
path.replace("/", "\\")
33+
return self.prepath + path
34+
1835
def connect(self) -> None:
1936
if self.connected:
2037
return
2138
try:
22-
self.ctx = SMBConnection(
23-
self.username,
24-
self.password,
25-
"smbclient",
26-
self.server,
27-
use_ntlm_v2=True,
39+
smbclient.register_session(
40+
self.server, port=self.port, **self.client_params
2841
)
29-
self.ctx.connect(self.server)
3042
self.connected = True
31-
except base.SMBTimeout as error:
43+
except SMBException as error:
3244
raise IOError(f"failed to connect: {error}")
3345

3446
def disconnect(self) -> None:
3547
self.connected = False
36-
try:
37-
self.ctx.close()
38-
except base.SMBTimeout as error:
39-
raise TimeoutError(f"disconnect: {error}")
48+
smbclient.reset_connection_cache(
49+
connection_cache=self.client_params["connection_cache"]
50+
)
51+
52+
def _check_connected(self, action: str) -> None:
53+
if not self.connected:
54+
raise ConnectionError(f"{action}: server not connected")
4055

4156
def listdir(self, path: str = "/") -> typing.List[str]:
57+
self._check_connected("listdir")
4258
try:
43-
dentries = self.ctx.listPath(self.share, path)
44-
except smb_structs.OperationFailure as error:
45-
raise IOError(f"failed to readdir: {error}")
46-
except base.SMBTimeout as error:
47-
raise TimeoutError(f"listdir: {error}")
48-
except base.NotConnectedError as error:
49-
raise ConnectionError(f"listdir: {error}")
50-
51-
return [dent.filename for dent in dentries]
59+
filenames = smbclient.listdir(
60+
self._path(path), **self.client_params
61+
)
62+
except SMBException as error:
63+
raise IOError(f"listdir: {error}")
64+
return filenames
5265

5366
def mkdir(self, dpath: str) -> None:
67+
self._check_connected("mkdir")
68+
if not self.connected:
69+
raise ConnectionError("listdir: server not connected")
5470
try:
55-
self.ctx.createDirectory(self.share, dpath)
56-
except smb_structs.OperationFailure as error:
57-
raise IOError(f"failed to mkdir: {error}")
58-
except base.SMBTimeout as error:
59-
raise TimeoutError(f"mkdir: {error}")
60-
except base.NotConnectedError as error:
61-
raise ConnectionError(f"mkdir: {error}")
71+
smbclient.mkdir(self._path(dpath), **self.client_params)
72+
except SMBException as error:
73+
raise IOError(f"mkdir: {error}")
6274

6375
def rmdir(self, dpath: str) -> None:
76+
self._check_connected("rmdir")
6477
try:
65-
self.ctx.deleteDirectory(self.share, dpath)
66-
except smb_structs.OperationFailure as error:
67-
raise IOError(f"failed to rmdir: {error}")
68-
except base.SMBTimeout as error:
69-
raise TimeoutError(f"rmdir: {error}")
70-
except base.NotConnectedError as error:
71-
raise ConnectionError(f"rmdir: {error}")
78+
smbclient.rmdir(self._path(dpath), **self.client_params)
79+
except SMBException as error:
80+
raise IOError(f"rmdir: {error}")
7281

7382
def unlink(self, fpath: str) -> None:
83+
self._check_connected("unlink")
7484
try:
75-
self.ctx.deleteFiles(self.share, fpath)
76-
except smb_structs.OperationFailure as error:
77-
raise IOError(f"failed to unlink: {error}")
78-
except base.SMBTimeout as error:
79-
raise TimeoutError(f"unlink: {error}")
80-
except base.NotConnectedError as error:
81-
raise ConnectionError(f"unlink: {error}")
85+
smbclient.remove(self._path(fpath), **self.client_params)
86+
except SMBException as error:
87+
raise IOError(f"unlink: {error}")
88+
89+
def _read_write_fd(self, fd_from: typing.IO, fd_to: typing.IO) -> None:
90+
while True:
91+
data = fd_from.read(rw_chunk_size)
92+
if not data:
93+
break
94+
n = 0
95+
while n < len(data):
96+
n += fd_to.write(data[n:])
8297

8398
def write(self, fpath: str, writeobj: typing.IO) -> None:
99+
self._check_connected("write")
84100
try:
85-
self.ctx.storeFile(self.share, fpath, writeobj)
86-
except smb_structs.OperationFailure as error:
87-
raise IOError(f"failed in write_text: {error}")
88-
except base.SMBTimeout as error:
89-
raise TimeoutError(f"write_text: {error}")
90-
except base.NotConnectedError as error:
91-
raise ConnectionError(f"write: {error}")
101+
with smbclient.open_file(
102+
self._path(fpath), mode="wb", **self.client_params
103+
) as fd:
104+
self._read_write_fd(writeobj, fd)
105+
except SMBException as error:
106+
raise IOError(f"write: {error}")
92107

93108
def read(self, fpath: str, readobj: typing.IO) -> None:
109+
self._check_connected("read")
94110
try:
95-
self.ctx.retrieveFile(self.share, fpath, readobj)
96-
except smb_structs.OperationFailure as error:
97-
raise IOError(f"failed in read_text: {error}")
98-
except base.SMBTimeout as error:
99-
raise TimeoutError(f"read_text: {error}")
100-
except base.NotConnectedError as error:
101-
raise ConnectionError(f"read: {error}")
111+
with smbclient.open_file(
112+
self._path(fpath), mode="rb", **self.client_params
113+
) as fd:
114+
self._read_write_fd(fd, readobj)
115+
except SMBException as error:
116+
raise IOError(f"write: {error}")
102117

103118
def write_text(self, fpath: str, teststr: str) -> None:
104-
with io.BytesIO(teststr.encode()) as writeobj:
105-
self.write(fpath, writeobj)
119+
self._check_connected("write_text")
120+
try:
121+
with smbclient.open_file(
122+
self._path(fpath), mode="w", **self.client_params
123+
) as fd:
124+
fd.write(teststr)
125+
except SMBException as error:
126+
raise IOError(f"write: {error}")
106127

107128
def read_text(self, fpath: str) -> str:
108-
with io.BytesIO() as readobj:
109-
self.read(fpath, readobj)
110-
ret = readobj.getvalue().decode("utf8")
129+
self._check_connected("read_text")
130+
try:
131+
with smbclient.open_file(
132+
self._path(fpath), **self.client_params
133+
) as fd:
134+
ret = fd.read()
135+
except SMBException as error:
136+
raise IOError(f"write: {error}")
111137
return ret

‎tox.ini

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ deps =
1313
pyyaml
1414
pytest-randomly
1515
iso8601
16-
pysmb
16+
smbprotocol
1717
commands = pytest -vrfEsxXpP testcases/
1818

1919
[testenv:pytest-unprivileged]
@@ -22,7 +22,7 @@ deps =
2222
pyyaml
2323
pytest-randomly
2424
iso8601
25-
pysmb
25+
smbprotocol
2626
commands = pytest -vrfEsxXpP -k 'not privileged' testcases/
2727

2828
[testenv:sanity]
@@ -31,7 +31,7 @@ deps =
3131
pyyaml
3232
pytest-randomly
3333
iso8601
34-
pysmb
34+
smbprotocol
3535
changedir = {toxinidir}
3636
commands = pytest -vrfEsxXpP testcases/consistency
3737

0 commit comments

Comments
 (0)