fixed chunk output

This commit is contained in:
Dionysus 2025-02-12 00:32:28 -05:00
parent db9590f59d
commit 41d7e53d30
Signed by: acidvegas
GPG Key ID: EF4B922DB85DC9DE
3 changed files with 18 additions and 26 deletions

View File

@ -6,4 +6,4 @@ from .colors import Colors
from .scanner import HTTPZScanner from .scanner import HTTPZScanner
__version__ = '2.1.1' __version__ = '2.1.2'

View File

@ -197,7 +197,7 @@ class HTTPZScanner:
self.resolvers = await load_resolvers(self.resolver_file) self.resolvers = await load_resolvers(self.resolver_file)
async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False)) as session: async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False)) as session:
tasks = set() tasks = {} # Change to dict to track domain for each task
domain_queue = asyncio.Queue() domain_queue = asyncio.Queue()
queue_empty = False queue_empty = False
@ -249,52 +249,44 @@ class HTTPZScanner:
try: try:
while not queue_empty or tasks: while not queue_empty or tasks:
# Fill up tasks to concurrent_limit # Start new tasks if needed
while len(tasks) < self.concurrent_limit and not queue_empty: while len(tasks) < self.concurrent_limit and not queue_empty:
try: try:
domain = await domain_queue.get() domain = await domain_queue.get()
if domain is None: # Queue is empty if domain is None:
queue_empty = True queue_empty = True
break break
task = asyncio.create_task(process_domain(domain)) task = asyncio.create_task(process_domain(domain))
tasks.add(task) tasks[task] = domain
except asyncio.CancelledError:
break
except Exception as e: except Exception as e:
debug(f'Error creating task: {str(e)}') debug(f'Error creating task: {str(e)}')
if not tasks: if not tasks:
break break
# Wait for any task to complete with timeout # Wait for the FIRST task to complete
try: try:
done, pending = await asyncio.wait( done, _ = await asyncio.wait(
tasks, tasks.keys(),
timeout=self.timeout, timeout=self.timeout,
return_when=asyncio.FIRST_COMPLETED return_when=asyncio.FIRST_COMPLETED
) )
# Handle completed tasks # Process completed task immediately
for task in done: for task in done:
tasks.remove(task) domain = tasks.pop(task)
try: try:
if result := await task: if result := await task:
yield result yield result
except Exception as e: except Exception as e:
debug(f'Error processing task result: {str(e)}') debug(f'Error processing result for {domain}: {str(e)}')
# Handle timed out tasks
if not done and pending:
for task in pending:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
tasks.remove(task)
except Exception as e: except Exception as e:
debug(f'Error in task processing loop: {str(e)}') debug(f'Error in task processing loop: {str(e)}')
# Remove any failed tasks
failed_tasks = [t for t in tasks if t.done() and t.exception()]
for task in failed_tasks:
tasks.pop(task)
finally: finally:
# Clean up # Clean up

View File

@ -10,7 +10,7 @@ with open('README.md', 'r', encoding='utf-8') as f:
setup( setup(
name='httpz_scanner', name='httpz_scanner',
version='2.1.1', version='2.1.2',
author='acidvegas', author='acidvegas',
author_email='acid.vegas@acid.vegas', author_email='acid.vegas@acid.vegas',
description='Hyper-fast HTTP Scraping Tool', description='Hyper-fast HTTP Scraping Tool',