Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import base64 

2import errno 

3import logging 

4import os 

5import re 

6import stat 

7from io import StringIO 

8 

9import paramiko 

10 

11from sdc_etl_libs.sdc_filetransfer.FileTransfer import FileTransfer 

12from sdc_etl_libs.sdc_credentials.sdc_endpoint_credentials import SDCEndpointCredentials 

13 

14class SFTP(FileTransfer): 

15 

16 def __init__(self): 

17 

18 self.host = None 

19 self.username = None 

20 self.port = None 

21 self.transport = None 

22 self.client = None 

23 self.rsa_key = None 

24 

25 def connect(self, host_, username_, port_, password_=None, rsa_key_=None): 

26 """ 

27 Connects to SFTP site and returns a handle to self.client. 

28 :param host_: SFTP hostname. 

29 :param username_: SFTP username. 

30 :param password_: SFTP password. Optional. 

31 :param rsa_key_: RSA Public or Private Key. Optional. 

32 :param port_: Int. SFTP port. 

33 :return: None. 

34 """ 

35 

36 self.host = host_ 

37 self.username = username_ 

38 self.port = int(port_) 

39 self.transport = paramiko.Transport((self.host, self.port)) 

40 

41 pkey = rsa_key_ 

42 

43 # Connect to SFTP site with rsa_key AND password 

44 if password_ and pkey: 

45 keyfile = StringIO(pkey) 

46 decrypt_key = paramiko.RSAKey.from_private_key(keyfile, password=password_) 

47 self.transport.start_client() 

48 self.transport.auth_publickey(username_, decrypt_key) 

49 

50 # Connect to SFTP site with rsa_key AND NOT password 

51 elif pkey and not password_: 

52 pkey = paramiko.RSAKey(file_obj=StringIO(pkey)) 

53 self.transport.start_client() 

54 self.transport.auth_publickey(username_, pkey) 

55 

56 # Connect to SFTP site with password AND NOT rsa_key 

57 elif password_ and not pkey: 

58 self.transport.connect(username=self.username, password=password_) 

59 

60 else: 

61 raise (Exception("Missing password or rsa_key in SFTPFIleTransfer.")) 

62 

63 self.client = paramiko.SFTPClient.from_transport(self.transport) 

64 

65 def check_if_path_exists(self, path_): 

66 """ 

67 Checks to see if path exists on SFTP site. 

68 :param path_: Path to object (Can be file for directory). 

69 :return: Boolean. 

70 """ 

71 

72 try: 

73 self.client.stat(path_) 

74 return True 

75 except IOError as e: 

76 if e.errno == errno.ENOENT: 

77 return False 

78 

79 def get_obj_list(self, path_=None, obj_regex_=None, give_full_path_=False, include_dirs_=True): 

80 """ 

81 Returns a list of objects in a path on an SFTP site. 

82 :param path_: Path to return object names from. 

83 Default '' (i.e. current directory). 

84 :param obj_regex_: Regex expression to search/return object 

85 names by. Default None. 

86 :param give_full_path_: If False, returns only object names. If True, 

87 If True, returns full path (from path_) of object. Default False. 

88 :param include_dirs_: If True, will remove directories 

89 from returned list results. If False, directories will be included. 

90 :return: List of SFTP object names as strings. Default False. 

91 """ 

92 

93 # If no path specified, set path to connection's starting directory. 

94 path = os.path.join(path_, '') if path_ else '' 

95 

96 if self.check_if_path_exists(path): 

97 

98 available_objects = [] 

99 object_results = [] 

100 

101 if include_dirs_: 

102 available_objects = self.client.listdir(path) 

103 else: 

104 for obj in self.client.listdir(path): 

105 full_path = os.path.join(path, obj) 

106 if not stat.S_ISDIR(self.get_obj_stats(full_path).st_mode): 

107 available_objects.append(obj) 

108 

109 if obj_regex_: 

110 available_objects = \ 

111 [x for x in available_objects if re.search(obj_regex_, x)] 

112 

113 for obj in available_objects: 

114 if give_full_path_: 

115 object_results.append(os.path.join(path, obj)) 

116 else: 

117 object_results.append(obj) 

118 return object_results 

119 

120 else: 

121 raise Exception(f"Path {path} does not exist on site.") 

122 

123 def get_obj_stats(self, obj_path_): 

124 """ 

125 Returns information about object on SFTP site. 

126 :param obj_path_: Path of object. 

127 :return: SFTPAttributes object that holds attributes returned by the 

128 site, which can vary. Attributes supported are st_mode, st_size, 

129 st_uid, st_gid, st_atime, st_mtime. 

130 """ 

131 

132 try: 

133 return self.client.stat(obj_path_) 

134 

135 except FileNotFoundError: 

136 logging.warning(f"{obj_path_} was not found on the SFTP site.") 

137 

138 except Exception as e: 

139 logging.warning(f"An error occurred. {e}") 

140 

141 def close_connection(self): 

142 """ 

143 Closes the SFTP connection. 

144 :return: None. 

145 """ 

146 

147 try: 

148 logging.info("Closing connection...") 

149 self.transport.close() 

150 logging.info("Connection closed.") 

151 

152 except Exception as e: 

153 logging.warning(f"There was an issue closing the connection: {e}") 

154 

155 def get_file_obj_record_count(self, file_obj_): 

156 """ 

157 Gets the total number of records contained in a file object. Returns file obj seek to 0 after record count. 

158 :param file_obj_: Object. File as an object from SFTP. 

159 :return: Int. Number of records in file object. 

160 """ 

161 

162 record_count = 0 

163 for line in file_obj_: 

164 record_count += 1 

165 file_obj_.seek(0) 

166 

167 return record_count 

168 

169 def get_file_as_file_object(self, file_path_): 

170 """ 

171 Returns a file as a file object. 

172 :param file_path_: Path of file on SFTP site. 

173 :return: File object, type of "paramiko.sftp_file.SFTPFile". 

174 """ 

175 

176 file_obj = self.client.file(file_path_, 'r') 

177 file_obj.seek(0) 

178 

179 return file_obj 

180 

181 def write_file(self, path_, file_name_, file_obj_, return_stats_=True): 

182 """ 

183 Writes a file object to SFTP site. 

184 :param path_: Path to write object to. 

185 :param file_name_: Name to write object as. 

186 :param file_obj_: File as an object to write out. 

187 :param return_stats_. Boolean. If True, checks item in S3 and returns size as confirmation of upload. 

188 :return: None. 

189 """ 

190 file_destination = os.path.join(path_, file_name_) 

191 

192 try: 

193 with self.client.open(file_destination, "w") as f: 

194 f.write(file_obj_.getvalue()) 

195 logging.info("File written to SFTP.") 

196 

197 except Exception as e: 

198 logging.exception(f"There was an issue writing this dataframe to the SFTP site: " f"{str(e)}") 

199 

200 if return_stats_: 

201 stats = self.get_obj_stats(file_destination) 

202 result = f"{round(stats.st_size / 1000000, 5):,} MB" 

203 return result 

204 else: 

205 return None