mirror of
https://github.com/romanz/amodem.git
synced 2026-02-24 16:18:12 +08:00
agent: add --shell and --connect flags
--shell: opens $SHELL as subprocess --connect: connect via SSH to specified identity's host
This commit is contained in:
@@ -13,43 +13,51 @@ log = logging.getLogger(__name__)
|
||||
def main():
|
||||
fmt = '%(asctime)s %(levelname)-12s %(message)-100s [%(filename)s:%(lineno)d]'
|
||||
p = argparse.ArgumentParser()
|
||||
g = p.add_mutually_exclusive_group()
|
||||
g.add_argument('-v', '--verbose', default=0, action='count')
|
||||
g.add_argument('-q', '--quiet', default=False, action='store_true')
|
||||
|
||||
p.add_argument('-p', '--public-key', default=False, action='store_true')
|
||||
p.add_argument('-v', '--verbose', default=0, action='count')
|
||||
|
||||
p.add_argument('-i', '--identity', type=str,
|
||||
help='proto://[user@]host[:port][/path]')
|
||||
p.add_argument('command', type=str, nargs='*',
|
||||
|
||||
g = p.add_mutually_exclusive_group()
|
||||
g.add_argument('-s', '--shell', default=False, action='store_true',
|
||||
help='run $SHELL as subprocess under SSH agent')
|
||||
g.add_argument('-c', '--connect', default=False, action='store_true',
|
||||
help='connect to specified host via SSH')
|
||||
p.add_argument('argument', type=str, nargs='*', metavar='ARGUMENT',
|
||||
help='command to run under the SSH agent')
|
||||
args = p.parse_args()
|
||||
|
||||
loglevel = logging.INFO
|
||||
if args.verbose:
|
||||
loglevel = logging.DEBUG
|
||||
if args.quiet:
|
||||
loglevel = logging.WARNING
|
||||
|
||||
logging.basicConfig(level=loglevel, format=fmt)
|
||||
levels = [logging.WARNING, logging.INFO, logging.DEBUG]
|
||||
level = levels[min(args.verbose, len(levels))]
|
||||
logging.basicConfig(level=level, format=fmt)
|
||||
|
||||
with trezor.Client(factory=trezor.TrezorLibrary) as client:
|
||||
identity = client.get_identity(label=args.identity)
|
||||
public_key = client.get_public_key(identity=identity)
|
||||
if args.public_key:
|
||||
|
||||
command, use_shell = args.argument, False
|
||||
if args.connect:
|
||||
to_ascii = lambda s: s.encode('ascii')
|
||||
command = ['ssh', to_ascii(identity.host)]
|
||||
if identity.user:
|
||||
command += ['-l', to_ascii(identity.user)]
|
||||
if identity.port:
|
||||
command += ['-p', to_ascii(identity.port)]
|
||||
log.debug('SSH connect: %r', command)
|
||||
|
||||
if args.shell:
|
||||
command, use_shell = os.environ['SHELL'], True
|
||||
log.debug('using shell: %r', command)
|
||||
|
||||
if not command:
|
||||
sys.stdout.write(public_key)
|
||||
return
|
||||
|
||||
command, use_shell = args.command, False
|
||||
if not command:
|
||||
command, use_shell = os.environ['SHELL'], True
|
||||
|
||||
signer = client.sign_ssh_challenge
|
||||
|
||||
try:
|
||||
with server.serve(public_keys=[public_key], signer=signer) as env:
|
||||
return server.run_process(
|
||||
command=command, environ=env, use_shell=use_shell
|
||||
)
|
||||
return server.run_process(command=command, environ=env,
|
||||
use_shell=use_shell)
|
||||
except KeyboardInterrupt:
|
||||
log.info('server stopped')
|
||||
|
||||
Reference in New Issue
Block a user