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