1# -*- coding: utf-8 -*-
2
3import os, tempfile
4from StringIO import StringIO
5from smb.SMBConnection import SMBConnection
6from util import getConnectionInfo
7from nose.tools import with_setup
8from smb import smb_structs
9
10try:
11    import hashlib
12    def MD5(): return hashlib.md5()
13except ImportError:
14    import md5
15    def MD5(): return md5.new()
16
17conn = None
18
19def setup_func_SMB1():
20    global conn
21    smb_structs.SUPPORT_SMB2 = False
22    info = getConnectionInfo()
23    conn = SMBConnection(info['user'], info['password'], info['client_name'], info['server_name'], use_ntlm_v2 = True, is_direct_tcp = True)
24    assert conn.connect(info['server_ip'], info['server_port'])
25
26def setup_func_SMB2():
27    global conn
28    smb_structs.SUPPORT_SMB2 = True
29    info = getConnectionInfo()
30    conn = SMBConnection(info['user'], info['password'], info['client_name'], info['server_name'], use_ntlm_v2 = True, is_direct_tcp = True)
31    assert conn.connect(info['server_ip'], info['server_port'])
32
33def teardown_func():
34    global conn
35    conn.close()
36
37@with_setup(setup_func_SMB1, teardown_func)
38def test_retr_multiplereads_SMB1():
39    # Test file retrieval using multiple ReadAndx calls (assuming each call will not reach more than 65534 bytes)
40    global conn
41    temp_fh = StringIO()
42    file_attributes, filesize = conn.retrieveFile('smbtest', '/rfc1001.txt', temp_fh)
43
44    md = MD5()
45    md.update(temp_fh.getvalue())
46    assert md.hexdigest() == '5367c2bbf97f521059c78eab65309ad3'
47    assert filesize == 158437
48
49    temp_fh.close()
50
51@with_setup(setup_func_SMB2, teardown_func)
52def test_retr_multiplereads_SMB2():
53    # Test file retrieval using multiple ReadAndx calls (assuming each call will not reach more than 65534 bytes)
54    global conn
55    temp_fh = StringIO()
56    file_attributes, filesize = conn.retrieveFile('smbtest', '/rfc1001.txt', temp_fh)
57
58    md = MD5()
59    md.update(temp_fh.getvalue())
60    assert md.hexdigest() == '5367c2bbf97f521059c78eab65309ad3'
61    assert filesize == 158437
62
63    temp_fh.close()
64
65@with_setup(setup_func_SMB1, teardown_func)
66def test_retr_longfilename_SMB1():
67    # Test file retrieval that has a long English filename
68    global conn
69    temp_fh = StringIO()
70    file_attributes, filesize = conn.retrieveFile('smbtest', '/Implementing CIFS - SMB.html', temp_fh)
71
72    md = MD5()
73    md.update(temp_fh.getvalue())
74    assert md.hexdigest() == '671c5700d279fcbbf958c1bba3c2639e'
75    assert filesize == 421269
76
77    temp_fh.close()
78
79@with_setup(setup_func_SMB2, teardown_func)
80def test_retr_longfilename_SMB2():
81    # Test file retrieval that has a long English filename
82    global conn
83    temp_fh = StringIO()
84    file_attributes, filesize = conn.retrieveFile('smbtest', '/Implementing CIFS - SMB.html', temp_fh)
85
86    md = MD5()
87    md.update(temp_fh.getvalue())
88    assert md.hexdigest() == '671c5700d279fcbbf958c1bba3c2639e'
89    assert filesize == 421269
90
91    temp_fh.close()
92
93@with_setup(setup_func_SMB1, teardown_func)
94def test_retr_unicodefilename_SMB1():
95    # Test file retrieval that has a long non-English filename inside a folder with a non-English name
96    global conn
97    temp_fh = StringIO()
98    file_attributes, filesize = conn.retrieveFile('smbtest', u'/测试文件夹/垃圾文件.dat', temp_fh)
99
100    md = MD5()
101    md.update(temp_fh.getvalue())
102    assert md.hexdigest() == '8a44c1e80d55e91c92350955cdf83442'
103    assert filesize == 256000
104
105    temp_fh.close()
106
107@with_setup(setup_func_SMB2, teardown_func)
108def test_retr_unicodefilename_SMB2():
109    # Test file retrieval that has a long non-English filename inside a folder with a non-English name
110    global conn
111    temp_fh = StringIO()
112    file_attributes, filesize = conn.retrieveFile('smbtest', u'/测试文件夹/垃圾文件.dat', temp_fh)
113
114    md = MD5()
115    md.update(temp_fh.getvalue())
116    assert md.hexdigest() == '8a44c1e80d55e91c92350955cdf83442'
117    assert filesize == 256000
118
119    temp_fh.close()
120
121@with_setup(setup_func_SMB1, teardown_func)
122def test_retr_offset_SMB1():
123    # Test file retrieval from offset to EOF
124    global conn
125    temp_fh = StringIO()
126    file_attributes, filesize = conn.retrieveFileFromOffset('smbtest', u'/测试文件夹/垃圾文件.dat', temp_fh, offset = 100000)
127
128    md = MD5()
129    md.update(temp_fh.getvalue())
130    assert md.hexdigest() == 'a141bd8024571ce7cb5c67f2b0d8ea0b'
131    assert filesize == 156000
132
133    temp_fh.close()
134
135@with_setup(setup_func_SMB2, teardown_func)
136def test_retr_offset_SMB2():
137    # Test file retrieval from offset to EOF
138    global conn
139    temp_fh = StringIO()
140    file_attributes, filesize = conn.retrieveFileFromOffset('smbtest', u'/测试文件夹/垃圾文件.dat', temp_fh, offset = 100000)
141
142    md = MD5()
143    md.update(temp_fh.getvalue())
144    assert md.hexdigest() == 'a141bd8024571ce7cb5c67f2b0d8ea0b'
145    assert filesize == 156000
146
147    temp_fh.close()
148
149@with_setup(setup_func_SMB1, teardown_func)
150def test_retr_offset_and_biglimit_SMB1():
151    # Test file retrieval from offset with a big max_length
152    global conn
153    temp_fh = StringIO()
154    file_attributes, filesize = conn.retrieveFileFromOffset('smbtest', u'/测试文件夹/垃圾文件.dat', temp_fh, offset = 100000, max_length = 100000)
155
156    md = MD5()
157    md.update(temp_fh.getvalue())
158    assert md.hexdigest() == '83b7afd7c92cdece3975338b5ca0b1c5'
159    assert filesize == 100000
160
161    temp_fh.close()
162
163@with_setup(setup_func_SMB2, teardown_func)
164def test_retr_offset_and_biglimit_SMB2():
165    # Test file retrieval from offset with a big max_length
166    global conn
167    temp_fh = StringIO()
168    file_attributes, filesize = conn.retrieveFileFromOffset('smbtest', u'/测试文件夹/垃圾文件.dat', temp_fh, offset = 100000, max_length = 100000)
169
170    md = MD5()
171    md.update(temp_fh.getvalue())
172    assert md.hexdigest() == '83b7afd7c92cdece3975338b5ca0b1c5'
173    assert filesize == 100000
174
175    temp_fh.close()
176
177@with_setup(setup_func_SMB1, teardown_func)
178def test_retr_offset_and_smalllimit_SMB1():
179    # Test file retrieval from offset with a small max_length
180    global conn
181    temp_fh = StringIO()
182    file_attributes, filesize = conn.retrieveFileFromOffset('smbtest', u'/测试文件夹/垃圾文件.dat', temp_fh, offset = 100000, max_length = 10)
183
184    md = MD5()
185    md.update(temp_fh.getvalue())
186    assert md.hexdigest() == '746f60a96b39b712a7b6e17ddde19986'
187    assert filesize == 10
188
189    temp_fh.close()
190
191@with_setup(setup_func_SMB2, teardown_func)
192def test_retr_offset_and_smalllimit_SMB2():
193    # Test file retrieval from offset with a small max_length
194    global conn
195    temp_fh = StringIO()
196    file_attributes, filesize = conn.retrieveFileFromOffset('smbtest', u'/测试文件夹/垃圾文件.dat', temp_fh, offset = 100000, max_length = 10)
197
198    md = MD5()
199    md.update(temp_fh.getvalue())
200    assert md.hexdigest() == '746f60a96b39b712a7b6e17ddde19986'
201    assert filesize == 10
202
203    temp_fh.close()
204
205@with_setup(setup_func_SMB1, teardown_func)
206def test_retr_offset_and_zerolimit_SMB1():
207    # Test file retrieval from offset to EOF with max_length=0
208    global conn
209    temp_fh = StringIO()
210    file_attributes, filesize = conn.retrieveFileFromOffset('smbtest', u'/测试文件夹/垃圾文件.dat', temp_fh, offset = 100000, max_length = 0)
211
212    md = MD5()
213    md.update(temp_fh.getvalue())
214    assert md.hexdigest() == 'd41d8cd98f00b204e9800998ecf8427e'
215    assert filesize == 0
216
217    temp_fh.close()
218
219@with_setup(setup_func_SMB2, teardown_func)
220def test_retr_offset_and_zerolimit_SMB2():
221    # Test file retrieval from offset to EOF with max_length=0
222    global conn
223    temp_fh = StringIO()
224    file_attributes, filesize = conn.retrieveFileFromOffset('smbtest', u'/测试文件夹/垃圾文件.dat', temp_fh, offset = 100000, max_length = 0)
225
226    md = MD5()
227    md.update(temp_fh.getvalue())
228    assert md.hexdigest() == 'd41d8cd98f00b204e9800998ecf8427e'
229    assert filesize == 0
230
231    temp_fh.close()
232