Fix call to net.WhoisClient

Lookyloo · Sep 7, 2016 · 3f16d42 · 3f16d42
1 parent 0d42862
commit 3f16d42
Showing 1 changed file with 12 additions and 12 deletions.
diff --git a/uwhoisd/__init__.py b/uwhoisd/__init__.py
@@ -160,18 +160,18 @@ def _run_query(self, server, port, query, prefix='', is_recursive=False):
                 logger.info("Rate limiting on %s", server)
                 self.redis_ratelimit.zremrangebyscore(max_key, '-inf', time.time())
                 time.sleep(1)
-        client = net.WhoisClient(server, port)
-        if is_recursive:
-            logger.info("Recursive query to %s about %s", server, query)
-        else:
-            logger.info("Querying %s about %s", server, query)
-        if self.redis_ratelimit is not None and ratelimit_details is not None:
-            self.redis_ratelimit.zremrangebyscore(max_key, '-inf', time.time())
-            self.redis_ratelimit.setex(server, ratelimit_details.split()[0], '')
-            self.redis_ratelimit.zadd(max_key, time.time() + 3600, query)
-        if prefix is not None:
-            query = '{} {}'.format(prefix, query)
-        return client.whois(query)
+        with net.WhoisClient(server, port) as client:
+            if is_recursive:
+                logger.info("Recursive query to %s about %s", server, query)
+            else:
+                logger.info("Querying %s about %s", server, query)
+            if self.redis_ratelimit is not None and ratelimit_details is not None:
+                self.redis_ratelimit.zremrangebyscore(max_key, '-inf', time.time())
+                self.redis_ratelimit.setex(server, ratelimit_details.split()[0], '')
+                self.redis_ratelimit.zadd(max_key, time.time() + 3600, query)
+            if prefix is not None:
+                query = '{} {}'.format(prefix, query)
+            return client.whois(query)
 
     def _thin_query(self, pattern, response, port, query):
         """