/
ssh.py
149 lines (119 loc) · 4.21 KB
/
ssh.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
# -*- coding: utf-8 -*-
#
# Copyright (C) 2019 Radim Rehurek <me@radimrehurek.com>
#
# This code is distributed under the terms and conditions
# from the MIT License (MIT).
#
"""Implements I/O streams over SSH.
Examples
--------
>>> with open('/proc/version_signature', host='1.2.3.4') as conn:
... print(conn.read())
b'Ubuntu 4.4.0-1061.70-aws 4.4.131'
Similarly, from a command line::
$ python -c "from smart_open import ssh;print(ssh.open('/proc/version_signature', host='1.2.3.4').read())"
b'Ubuntu 4.4.0-1061.70-aws 4.4.131'
"""
import getpass
import logging
import urllib.parse
import warnings
import smart_open.utils
logger = logging.getLogger(__name__)
#
# Global storage for SSH connections.
#
_SSH = {}
SCHEMES = ("ssh", "scp", "sftp")
"""Supported URL schemes."""
DEFAULT_PORT = 22
URI_EXAMPLES = (
'ssh://username@host/path/file',
'ssh://username@host//path/file',
'scp://username@host/path/file',
'sftp://username@host/path/file',
)
def _unquote(text):
return text and urllib.parse.unquote(text)
def parse_uri(uri_as_string):
split_uri = urllib.parse.urlsplit(uri_as_string)
assert split_uri.scheme in SCHEMES
return dict(
scheme=split_uri.scheme,
uri_path=_unquote(split_uri.path),
user=_unquote(split_uri.username),
host=split_uri.hostname,
port=int(split_uri.port or DEFAULT_PORT),
password=_unquote(split_uri.password),
)
def open_uri(uri, mode, transport_params):
smart_open.utils.check_kwargs(open, transport_params)
parsed_uri = parse_uri(uri)
uri_path = parsed_uri.pop('uri_path')
parsed_uri.pop('scheme')
return open(uri_path, mode, transport_params=transport_params, **parsed_uri)
def _connect(hostname, username, port, password, transport_params):
try:
import paramiko
except ImportError:
warnings.warn(
'paramiko missing, opening SSH/SCP/SFTP paths will be disabled. '
'`pip install paramiko` to suppress'
)
raise
key = (hostname, username)
ssh = _SSH.get(key)
if ssh is None:
ssh = _SSH[key] = paramiko.client.SSHClient()
ssh.load_system_host_keys()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
kwargs = transport_params.get('connect_kwargs', {}).copy()
# if 'key_filename' is present in transport_params, then I do not
# overwrite the credentials.
if 'key_filename' not in kwargs:
kwargs.setdefault('password', password)
kwargs.setdefault('username', username)
ssh.connect(hostname, port, **kwargs)
return ssh
def open(path, mode='r', host=None, user=None, password=None, port=DEFAULT_PORT, transport_params=None):
"""Open a file on a remote machine over SSH.
Expects authentication to be already set up via existing keys on the local machine.
Parameters
----------
path: str
The path to the file to open on the remote machine.
mode: str, optional
The mode to use for opening the file.
host: str, optional
The hostname of the remote machine. May not be None.
user: str, optional
The username to use to login to the remote machine.
If None, defaults to the name of the current user.
password: str, optional
The password to use to login to the remote machine.
port: int, optional
The port to connect to.
transport_params: dict, optional
Any additional settings to be passed to paramiko.SSHClient.connect
Returns
-------
A file-like object.
Important
---------
If you specify a previously unseen host, then its host key will be added to
the local ~/.ssh/known_hosts *automatically*.
If ``username`` or ``password`` are specified in *both* the uri and
``transport_params``, ``transport_params`` will take precedence
"""
if not host:
raise ValueError('you must specify the host to connect to')
if not user:
user = getpass.getuser()
if not transport_params:
transport_params = {}
conn = _connect(host, user, port, password, transport_params)
sftp_client = conn.get_transport().open_sftp_client()
fobj = sftp_client.open(path, mode)
fobj.name = path
return fobj