1#!/usr/bin/env python
2
3# Copyright (C) 2011  Jeff Forcier <jeff@bitprophet.org>
4#
5# This file is part of ssh.
6#
7# 'ssh' is free software; you can redistribute it and/or modify it under the
8# terms of the GNU Lesser General Public License as published by the Free
9# Software Foundation; either version 2.1 of the License, or (at your option)
10# any later version.
11#
12# 'ssh' is distrubuted in the hope that it will be useful, but WITHOUT ANY
13# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
14# A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
15# details.
16#
17# You should have received a copy of the GNU Lesser General Public License
18# along with 'ssh'; if not, write to the Free Software Foundation, Inc.,
19# 51 Franklin Street, Suite 500, Boston, MA  02110-1335  USA.
20
21"""
22do the unit tests!
23"""
24
25import os
26import re
27import sys
28import unittest
29from optparse import OptionParser
30import ssh
31import threading
32
33sys.path.append('tests')
34
35from test_message import MessageTest
36from test_file import BufferedFileTest
37from test_buffered_pipe import BufferedPipeTest
38from test_util import UtilTest
39from test_hostkeys import HostKeysTest
40from test_pkey import KeyTest
41from test_kex import KexTest
42from test_packetizer import PacketizerTest
43from test_auth import AuthTest
44from test_transport import TransportTest
45from test_sftp import SFTPTest
46from test_sftp_big import BigSFTPTest
47from test_client import SSHClientTest
48
49default_host = 'localhost'
50default_user = os.environ.get('USER', 'nobody')
51default_keyfile = os.path.join(os.environ.get('HOME', '/'), '.ssh/id_rsa')
52default_passwd = None
53
54
55def iter_suite_tests(suite):
56    """Return all tests in a suite, recursing through nested suites"""
57    for item in suite._tests:
58        if isinstance(item, unittest.TestCase):
59            yield item
60        elif isinstance(item, unittest.TestSuite):
61            for r in iter_suite_tests(item):
62                yield r
63        else:
64            raise Exception('unknown object %r inside test suite %r'
65                            % (item, suite))
66
67
68def filter_suite_by_re(suite, pattern):
69    result = unittest.TestSuite()
70    filter_re = re.compile(pattern)
71    for test in iter_suite_tests(suite):
72        if filter_re.search(test.id()):
73            result.addTest(test)
74    return result
75
76
77def main():
78    parser = OptionParser('usage: %prog [options]')
79    parser.add_option('--verbose', action='store_true', dest='verbose', default=False,
80                      help='verbose display (one line per test)')
81    parser.add_option('--no-pkey', action='store_false', dest='use_pkey', default=True,
82                      help='skip RSA/DSS private key tests (which can take a while)')
83    parser.add_option('--no-transport', action='store_false', dest='use_transport', default=True,
84                      help='skip transport tests (which can take a while)')
85    parser.add_option('--no-sftp', action='store_false', dest='use_sftp', default=True,
86                      help='skip SFTP client/server tests, which can be slow')
87    parser.add_option('--no-big-file', action='store_false', dest='use_big_file', default=True,
88                      help='skip big file SFTP tests, which are slow as molasses')
89    parser.add_option('-R', action='store_false', dest='use_loopback_sftp', default=True,
90                      help='perform SFTP tests against a remote server (by default, SFTP tests ' +
91                      'are done through a loopback socket)')
92    parser.add_option('-H', '--sftp-host', dest='hostname', type='string', default=default_host,
93                      metavar='<host>',
94                      help='[with -R] host for remote sftp tests (default: %s)' % default_host)
95    parser.add_option('-U', '--sftp-user', dest='username', type='string', default=default_user,
96                      metavar='<username>',
97                      help='[with -R] username for remote sftp tests (default: %s)' % default_user)
98    parser.add_option('-K', '--sftp-key', dest='keyfile', type='string', default=default_keyfile,
99                      metavar='<keyfile>',
100                      help='[with -R] location of private key for remote sftp tests (default: %s)' %
101                      default_keyfile)
102    parser.add_option('-P', '--sftp-passwd', dest='password', type='string', default=default_passwd,
103                      metavar='<password>',
104                      help='[with -R] (optional) password to unlock the private key for remote sftp tests')
105
106    options, args = parser.parse_args()
107
108    # setup logging
109    ssh.util.log_to_file('test.log')
110
111    if options.use_sftp:
112        if options.use_loopback_sftp:
113            SFTPTest.init_loopback()
114        else:
115            SFTPTest.init(options.hostname, options.username, options.keyfile, options.password)
116        if not options.use_big_file:
117            SFTPTest.set_big_file_test(False)
118
119    suite = unittest.TestSuite()
120    suite.addTest(unittest.makeSuite(MessageTest))
121    suite.addTest(unittest.makeSuite(BufferedFileTest))
122    suite.addTest(unittest.makeSuite(BufferedPipeTest))
123    suite.addTest(unittest.makeSuite(UtilTest))
124    suite.addTest(unittest.makeSuite(HostKeysTest))
125    if options.use_pkey:
126        suite.addTest(unittest.makeSuite(KeyTest))
127    suite.addTest(unittest.makeSuite(KexTest))
128    suite.addTest(unittest.makeSuite(PacketizerTest))
129    if options.use_transport:
130        suite.addTest(unittest.makeSuite(AuthTest))
131        suite.addTest(unittest.makeSuite(TransportTest))
132    suite.addTest(unittest.makeSuite(SSHClientTest))
133    if options.use_sftp:
134        suite.addTest(unittest.makeSuite(SFTPTest))
135    if options.use_big_file:
136        suite.addTest(unittest.makeSuite(BigSFTPTest))
137    verbosity = 1
138    if options.verbose:
139        verbosity = 2
140
141    runner = unittest.TextTestRunner(verbosity=verbosity)
142    if len(args) > 0:
143        filter = '|'.join(args)
144        suite = filter_suite_by_re(suite, filter)
145    result = runner.run(suite)
146    # Clean up stale threads from poorly cleaned-up tests.
147    # TODO: make that not a problem, jeez
148    for thread in threading.enumerate():
149        if thread is not threading.currentThread():
150            thread._Thread__stop()
151    # Exit correctly
152    if not result.wasSuccessful():
153        sys.exit(1)
154
155
156if __name__ == '__main__':
157    main()
158