Test your AI Agents with the all-new Agent to Agent Testing Platform.Learn More

How to use runner_queue method in avocado

Best Python code snippet using avocado_python

main.py

Source:main.py

1from ..config_validate import load_conf_file2from ..util import Context3from ..db import build_database, Team, Challenge, ChallengeInst, FlagStatus, Flag, ExpStatus, ExpLog4from ..executor import ExpRunner, FlagSubmitter, pool_init, pool_stop, submit_init, submit_stop5from .time_watcher import alarm_start, alarm_stop, current_round, round_range6from .file_watcher import watcher_start, watcher_stop7from .exp_manager2 import ExpManager8from .exploit import ExploitScheduler9import signal10import threading11import queue12def testfunc():13    with open("config.yaml") as fp:14        load_conf_file(fp)15    build_database()16    Context.cv = threading.Condition()17    e = threading.Event()18    watcher_start(e)19    manager = ExpManager("pwn1")20    for i in range(4):21        exp = manager.exp_iter()22        if exp:23            print(exp.name)24        else:25            print("stop")26    print("wait")27    while True:28        try:29            with Context.cv:30                Context.cv.wait_for(lambda: e.is_set())31            e.clear()32            manager.check_exp_update()33            while True:34                exp = manager.exp_iter()35                if exp:36                    print(exp.name)37                else:38                    print("end")39                    break40        except KeyboardInterrupt:41            break42def signal_init(flag_submit, time_event):43    signal.signal(signal.SIGUSR1, lambda: flag_submit.notify_new_flag())44    signal.signal(signal.SIGUSR2, lambda: time_event.set())45class FlagSubmitQueue(object):46    def __init__(self, notify_queue):47        self.notify_queue = notify_queue48        self.free_worker_count = Context.max_submitters49        # TODO: not good50        if self.free_worker_count is None:51            self.free_worker_count = 10052        self.notify_new_flag()  # submit flag that not upload53    def notify_submit_flag(self):54        self.free_worker_count += 155    def notify_new_flag(self):56        if self.free_worker_count != 0:57            start, end = round_range(current_round())58            session = Context.db.get_session()59            flags = session.query(Flag).filter(Flag.timestamp >= start, Flag.timestamp <= end,60                                               Flag.submit_status == FlagStatus.wait_submit).order_by(61                Flag.weight).all()62            c = min(len(flags), self.free_worker_count)63            for i in range(c):64                f = flags[i]65                e = None66                for exploit in Context.exploit_list:67                    if exploit.inst_id == f.inst_id:68                        e = exploit69                        break70                submitter = FlagSubmitter(f.id, self.notify_queue, e)71                submitter.start()72            self.free_worker_count -= c73            session.commit()74            session.close()75def server_start(config_file: str):76    with open(config_file) as fp:77        load_conf_file(fp)78    build_database()79    Context.cv = threading.Condition()80    # start alarm clock81    time_event = threading.Event()82    time_event.set()83    alarm_start(time_event)84    # start file watch85    file_event = threading.Event()86    watcher_start(file_event)87    # start runner thread pool88    runner_queue = queue.Queue()89    # start submitter pool90    submitter_queue = queue.Queue()91    flag_submit_queue = FlagSubmitQueue(submitter_queue)92    def task_down():93        if not submitter_queue.empty():94            return True95        elif not runner_queue.empty():96            return True97        elif file_event.is_set():98            return True99        elif time_event.is_set():100            return True101        else:102            return False103    signal_init(flag_submit_queue, time_event)104    # the thread pool will init in the first submit os task105    while True:106        try:107            with Context.cv:108                Context.cv.wait_for(task_down)109            # a new round is start110            if time_event.is_set():111                Context.logger.info("new round")112                time_event.clear()113                session = Context.db.get_session()114                now_round = current_round()115                for exploit in Context.exploit_list:116                    exploit.stop()117                Context.exploit_list = []118                start, end = round_range(current_round())119                for flag in session.query(Flag).filter(Flag.timestamp < start,120                                                       Flag.submit_status == FlagStatus.wait_submit):121                    flag.submit_status = FlagStatus.flag_expire122                session.commit()123                for inst in session.query(ChallengeInst).join(Challenge).join(Team).filter(Team.active == 1,124                                                                                           Challenge.active == 1).all():125                    exploit = ExploitScheduler(inst.id, now_round, runner_queue, submitter_queue)126                    exploit.start()127                    Context.exploit_list.append(exploit)128                session.commit()129                session.close()130            if file_event.is_set():131                file_event.clear()132                for inst in Context.exploit_list:133                    inst.exp_modify()134            while not runner_queue.empty():135                log_id, exploit, is_success = runner_queue.get()136                exploit.task_update(log_id, is_success)137                if is_success:138                    flag_submit_queue.notify_new_flag()139            while not submitter_queue.empty():140                flag_id, exploit = submitter_queue.get()141                if exploit:142                    exploit.flag_submitted(flag_id)143                flag_submit_queue.notify_submit_flag()144        except KeyboardInterrupt:145            break146    # stop all thing and exit147    alarm_stop()148    pool_stop()149    submit_stop()...

exploit.py

Source:exploit.py

1from ..util import Context2from ..db import ChallengeInst, Flag, FlagStatus3from .exp_manager2 import ExpManager4from ..executor import ExpRunner, FlagSubmitter5class ExploitScheduler:6    def __init__(self, inst_id, round_count, runner_queue, flag_queue):7        self.inst_id = inst_id8        self.round = round_count9        self.runner_queue = runner_queue10        self.flag_queue = flag_queue11        # get exp and sorted it12        session = Context.db.get_session()13        inst = session.query(ChallengeInst).filter(ChallengeInst.id == inst_id).one()14        self.challenge_name = inst.challenge.name15        self.team_name = inst.team.name16        session.commit()17        session.close()18        self.exp_manage = ExpManager(self.challenge_name)19        self.stopped = False20        self.success = False21        self.task_wait = False22        self.run_all = False23    def run_new_one(self):24        # if this inst has a flag wait to submit, not run a new one25        session = Context.db.get_session()26        inst = session.query(ChallengeInst).filter(ChallengeInst.id == self.inst_id).one()27        for i in inst.flags:28            if i.submit_status == FlagStatus.wait_submit:29                return30        exp = self.exp_manage.exp_iter()31        if exp:32            runner = ExpRunner(exp.name, exp.entry, self.inst_id, self.runner_queue, self, exp.timeout)33            runner.start()34            self.task_wait = True35        else:36            self.run_all = True37    def start(self):38        self.run_new_one()39    def task_update(self, log_id, get_flag):40        self.task_wait = False41        if not get_flag and not self.stopped and not self.success:42            self.run_new_one()43        # if get flag or stopped, do not run new exploit44    def exp_modify(self):45        if self.stopped or self.success:46            # not need to update47            return48        self.exp_manage.check_exp_update()49    def flag_submitted(self, flag_id):50        session = Context.db.get_session()51        flag = session.query(Flag).filter(Flag.id == flag_id).one()52        if flag.submit_status == FlagStatus.submit_success:53            # not need run new task54            self.success = True55            return56        session.commit()57        session.close()58        # the task wait will be true when a duplicate flag submit failed59        if not self.stopped and not self.task_wait:60            self.run_new_one()61    def stop(self):62        self.stopped = True63    def renew(self):64        if self.run_all:65            self.exp_manage = ExpManager(self.challenge_name)...

streams.py

Source:streams.py

1"""Components that manage Dapp Runner's data and state streams."""2import asyncio3from collections import defaultdict4from dataclasses import dataclass5from typing import Callable, Dict, List, Optional, TextIO, Any, TypeVar, Generic6Msg = TypeVar("Msg")7@dataclass8class RunnerStream(Generic[Msg]):9    """Dapp Runner's output stream manager."""10    queue: asyncio.Queue11    stream: TextIO12    process_callback: Optional[Callable] = None13    """callback processing the queue messages"""14    async def update(self):15        """Await the queue and write to the output stream."""16        while True:17            try:18                msg = await self.queue.get()19            except asyncio.CancelledError:20                return21            if self.process_callback:22                msg = self.process_callback(msg)23            self.stream.write(str(msg) + "\n")24class RunnerStreamer:25    """Dapp Runner's stream writer."""26    _streams: Dict[asyncio.Queue, List[RunnerStream]]27    _tasks: List[asyncio.Task]28    def __init__(self):29        self._streams = defaultdict(list)30        self._tasks = []31    def register_stream(32        self,33        runner_queue: asyncio.Queue,34        stream: TextIO,35        process_callback: Optional[Callable] = None,36    ):37        """Register a stream and run the stream update task."""38        if runner_queue not in self._streams:39            self._init_queue(runner_queue)40        runner_stream: RunnerStream[Any] = RunnerStream(41            asyncio.Queue(), stream, process_callback42        )43        self._streams[runner_queue].append(runner_stream)44        self._tasks.append(asyncio.create_task(runner_stream.update()))45    async def _feed_queue(self, queue: asyncio.Queue):46        while True:47            try:48                msg = await queue.get()49            except asyncio.CancelledError:50                return51            for runner_stream in self._streams[queue]:52                runner_stream.queue.put_nowait(msg)53    def _init_queue(self, runner_queue: asyncio.Queue):54        """Start the feed task for the given queue."""55        self._tasks.append(asyncio.create_task(self._feed_queue(runner_queue)))56    async def stop(self):57        """Stop the stream feed tasks."""58        for t in self._tasks:59            t.cancel()...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.