"""
Copyright (c) [2021] Huawei Technologies Co.,Ltd.ALL rights reserved.
This program is licensed under Mulan PSL v2.
You can use it according to the terms and conditions of the Mulan PSL v2.
http://license.coscl.org.cn/MulanPSL2
THIS PROGRAM IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details.
@Author : lemon-higgins
@email : lemon.higgins@aliyun.com
@Date : 2021-04-21 16:14:27
@License : Mulan PSL v2
@Version : 1.0
@Desc : 文件传输
"""
import os
import sys
import stat
import re
import paramiko
import argparse
import subprocess
SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
sys.path.append(SCRIPT_PATH)
import mugen_log
import ssh_cmd
def get_remote_file(sftp, remote_dir, remote_file=None):
"""获取对端文件
Args:
sftp (class): 和对端建立连接
remote_dir ([str]): 远端需要传输的文件所在的目录
remote_file ([str], optional): 远端需要传输的文件. Defaults to None.
Returns:
[list]: 文件列表
"""
all_file = list()
remote_dir = remote_dir.rstrip("/")
dir_files = sftp.listdir_attr(remote_dir)
for d_f in dir_files:
if remote_file is not None and re.match(remote_file, d_f.filename) is None:
continue
_name = remote_dir + "/" + d_f.filename
if stat.S_ISDIR(d_f.st_mode):
all_file.extend(get_remote_file(sftp, _name))
else:
all_file.append(_name)
return all_file
def psftp_get(conn, remote_dir, remote_file="", local_dir=os.getcwd()):
"""获取远端文件
Args:
conn ([class]): 和远端建立连接
remote_dir ([str]): 远端需要传输的文件所在的目录
remote_file ([str], optional): 远端需要传输的文件. Defaults to None.
local_dir ([str], optional): 本地存放文件的目录. Defaults to os.getcwd().
"""
if conn == 519:
sys.exit(519)
sftp = paramiko.SFTPClient.from_transport(conn.get_transport())
if ssh_cmd.pssh_cmd(conn, "test -d " + remote_dir)[0]:
mugen_log.logging("error", "remote dir:%s does not exist" % remote_dir)
conn.close()
sys.exit(1)
all_file = list()
if remote_file == "":
all_file = get_remote_file(sftp, remote_dir)
else:
if ssh_cmd.pssh_cmd(conn, "test -f " + os.path.join(remote_dir, remote_file))[0]:
mugen_log.logging("error", "remote file:%s does not exist" % remote_file)
conn.close()
sys.exit(1)
all_file = get_remote_file(sftp, remote_dir, remote_file)
local_dir = os.path.normpath(local_dir)
remote_dir = os.path.normpath(remote_dir)
for f in all_file:
if remote_file == "":
storage_dir = remote_dir.split("/")[-1]
storage_path = os.path.join(
local_dir, storage_dir + os.path.dirname(f[len(remote_dir):])
)
if not os.path.exists(storage_path):
os.makedirs(storage_path)
sftp.get(f, os.path.join(storage_path, f.split("/")[-1]))
else:
if not os.path.exists(local_dir):
os.makedirs(local_dir)
sftp.get(f, os.path.join(local_dir, f.split("/")[-1]))
mugen_log.logging("info", "start to get file:%s......" % f)
conn.close()
def get_local_file(local_dir, local_file=None):
"""获取本地文件列表
Args:
local_dir ([str]): 本地文件所在的目录
local_file ([str], optional): 本地需要传输的文件. Defaults to None.
Returns:
[list]: 文件列表
"""
all_file = list()
local_dir = local_dir.rstrip("/")
dir_files = os.listdir(local_dir)
for d_f in dir_files:
if local_file is not None and re.match(local_file, d_f) is None:
continue
_name = local_dir + "/" + d_f
if os.path.isdir(_name):
all_file.extend(get_local_file(_name))
else:
all_file.append(_name)
return all_file
def psftp_put(conn, local_dir=os.getcwd(), local_file="", remote_dir=""):
"""将本地文件传输到远端
Args:
conn ([class]): 和远端建立连接
local_dir ([str]): 本地文件所在的目录
local_file ([str], optional): 本地需要传输的文件. Defaults to None.
remote_dir (str, optional): 远端存放文件的目录. Defaults to 根目录.
"""
if conn == 519:
sys.exit(519)
sftp = paramiko.SFTPClient.from_transport(conn.get_transport())
if subprocess.getstatusoutput("test -d " + local_dir)[0]:
mugen_log.logging("error", "local dir:%s does not exist" % local_dir)
conn.close()
sys.exit(1)
all_file = list()
if local_file == "":
all_file = get_local_file(local_dir)
else:
if subprocess.getstatusoutput("test -f " + os.path.join(local_dir, local_file))[0]:
mugen_log.logging("error", "local file:%s does not exist" % local_file)
conn.close()
sys.exit(1)
all_file = get_local_file(local_dir, local_file)
if remote_dir == "":
remote_dir = ssh_cmd.pssh_cmd(conn, "pwd")[1]
local_dir = os.path.normpath(local_dir)
remote_dir = os.path.normpath(remote_dir)
for f in all_file:
if local_file == "":
storage_dir = local_dir.split("/")[-1]
storage_path = os.path.join(
remote_dir, storage_dir + os.path.dirname(f[len(local_dir):])
)
if ssh_cmd.pssh_cmd(conn, "test -d " + storage_path)[0]:
ssh_cmd.pssh_cmd(conn, "mkdir -p " + storage_path)
sftp.put(f, os.path.join(storage_path, f.split("/")[-1]))
else:
if ssh_cmd.pssh_cmd(conn, "test -d " + remote_dir)[0]:
ssh_cmd.pssh_cmd(conn, "mkdir -p " + remote_dir)
sftp.put(f, os.path.join(remote_dir, f.split("/")[-1]))
mugen_log.logging("info", "start to put file:%s......" % f)
conn.close()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="manual to this script")
parser.add_argument("operation", type=str, choices=["get", "put"], default=None)
parser.add_argument(
"--remotedir", type=str, default=None, help="Must be an absolute path"
)
parser.add_argument("--node", type=int, default=2)
parser.add_argument("--remotefile", type=str, default="")
parser.add_argument("--localdir", type=str, default=os.getcwd())
parser.add_argument("--localfile", type=str, default="")
parser.add_argument("--ip", type=str, default=None)
parser.add_argument("--password", type=str, default=None)
parser.add_argument("--port", type=int, default=22)
parser.add_argument("--user", type=str, default="root")
parser.add_argument("--timeout", type=int, default=None)
args = parser.parse_args()
if args.node is not None and args.node > 0:
args.ip = os.environ.get("NODE" + str(args.node) + "_IPV4")
args.password = os.environ.get("NODE" + str(args.node) + "_PASSWORD")
args.port = os.environ.get("NODE" + str(args.node) + "_SSH_PORT")
args.user = os.environ.get("NODE" + str(args.node) + "_USER")
if (
args.ip is None
or args.password is None
or args.port is None
or args.user is None
):
mugen_log.logging(
"error",
"You need to check the environment configuration file to see if this node information exists.",
)
sys.exit(1)
if args.node < 0:
if args.ip is None or args.password is None:
mugen_log.logging(
"error",
"You need to check the parameter if node id < 0, ip and password must set.",
)
sys.exit(1)
conn = ssh_cmd.pssh_conn(args.ip, args.password, args.port, args.user, args.timeout)
if sys.argv[1] == "get":
psftp_get(conn, args.remotedir, args.remotefile, args.localdir)
elif sys.argv[1] == "put":
psftp_put(conn, args.localdir, args.localfile, args.remotedir)
else:
mugen_log.logging("error", "the following arguments are required:get|put")
sys.exit(1)