1 files changed, 12 insertions(+), 6 deletions(-)

M certtool/__main__.py

@@ 66,10 66,11 @@ def check_ocsp(domain, port, chain):
         print(e)
 
 
-def get_tls_certificate(domain, port):
+def get_tls_certificate(domain, port, timeout=5):
     ctx = OpenSSL.SSL.Context(OpenSSL.SSL.TLSv1_2_METHOD)
-    conn = socket.create_connection((domain, port))
+    conn = socket.create_connection((domain, port), timeout=timeout)
     sock = OpenSSL.SSL.Connection(ctx, conn)
+    sock.setblocking(True)
     sock.set_connect_state()
     sock.set_tlsext_host_name(domain.encode())
     sock.do_handshake()


@@ 77,8 78,11 @@ def get_tls_certificate(domain, port):
     return chain
 
 
-def check_generic(domain, port):
-    chain = get_tls_certificate(domain, port)
+def check_generic(domain, port, timeout=5):
+    try:
+        chain = get_tls_certificate(domain, port, timeout=timeout)
+    except TimeoutError:
+        return False, "Timeout"
     x509 = chain[0]
     components = x509.get_subject().get_components()
     for component in components:


@@ 154,6 158,8 @@ def main():
     parser.add_argument('--pop3s', action=argparse.BooleanOptionalAction, default=False, help='Check pop3s')
     parser.add_argument('--smtps', action=argparse.BooleanOptionalAction, default=False, help='Check smtp')
     parser.add_argument('--port', action="append", help="Check a specific port", default=[], type=int)
+    parser.add_argument('--timeout', '-t', default=5, help='Set the timeout in seconds for the TCP connection',
+                        type=int)
     args = parser.parse_args()
 
     domain_len = len(max(args.domain, key=len))


@@ 185,12 191,12 @@ def main():
             if len(protolist) > 1:
                 if not first_proto:
                     print(' ' * (domain_len + 2), end='')
-                print(f'{str(proto).ljust(5)} ', end='')
+                print(f'{str(proto).ljust(5)} ', end='', flush=True)
             first_proto = False
 
             if isinstance(proto, str):
                 proto = protocols[proto]
-            result = check_generic(domain, proto)
+            result = check_generic(domain, proto, timeout=args.timeout)
 
             if result[0] is True:
                 print(colorama.Fore.GREEN, result[1], colorama.Style.RESET_ALL)