root/branch/PacketManipulator/PM/Backend/Scapy/Context/Sequence.py @ 3672

Revision 3672, 6.1 kB (checked in by nopper, 5 years ago)

Implementing save and load in SequenceContext? (Scapy)
Line 
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3# Copyright (C) 2008 Adriano Monteiro Marques
4#
5# Author: Francesco Piccinno <stack.box@gmail.com>
6#
7# This program is free software; you can redistribute it and/or modify
8# it under the terms of the GNU General Public License as published by
9# the Free Software Foundation; either version 2 of the License, or
10# (at your option) any later version.
11#
12# This program is distributed in the hope that it will be useful,
13# but WITHOUT ANY WARRANTY; without even the implied warranty of
14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15# GNU General Public License for more details.
16#
17# You should have received a copy of the GNU General Public License
18# along with this program; if not, write to the Free Software
19# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
20
21from __future__ import with_statement
22from threading import Lock
23
24from PM.Core.I18N import _
25from PM.Backend.Scapy.serialize import load_sequence, save_sequence
26from PM.Backend.Scapy.utils import execute_sequence
27
28def register_sequence_context(BaseSequenceContext):
29   
30    class SequenceContext(BaseSequenceContext):
31        file_types = [(_('Scapy sequence'), '*.pms')]
32
33        def __init__(self, seq, count, inter, iface, strict, report_recv, report_sent,
34                     scallback, rcallback, sudata=None, rudata=None):
35
36            BaseSequenceContext.__init__(self, seq, count, inter, iface,
37                                         strict, report_recv, report_sent,
38                                         scallback, rcallback, sudata, rudata)
39
40            self.sequencer = None
41            self.internal = False
42            self.lock = Lock()
43            self.title = _('Unsaved')
44            self.summary = _('Executing sequence (%d packets %d times)') % (self.tot_packet_count, count)
45       
46        def save(self):
47            if self.cap_file:
48                self.seq = self.load_sequence(self.cap_file)
49
50                if self.seq is not None:
51                    self.state = self.SAVED
52                    return True
53
54            self.state = self.NOT_SAVED
55            return False
56
57        def load(self):
58            if self.cap_file and self.seq is not None and \
59               save_sequence(self.cap_file, self.seq):
60
61                self.state = self.SAVED
62                return True
63
64            self.state = self.NOT_SAVED
65            return False
66
67        def get_all_data(self):
68            with self.lock:
69                return BaseSequenceContext.get_all_data(self)
70
71        def get_data(self):
72            with self.lock:
73                return BaseSequenceContext.get_data(self)
74
75        def set_data(self, val):
76            with self.lock:
77                self.data = val
78
79        def _start(self):
80            if self.tot_packet_count - self.packet_count > 0 or \
81               self.tot_loop_count - self.loop_count > 0:
82
83                self.internal = True
84                self.state = self.RUNNING
85
86                self.sequencer = execute_sequence(
87                    self.seq,
88                    max(self.tot_loop_count - self.loop_count, 1),
89                    self.inter, self.iface, self.strict,
90                    self.__send_callback, self.__recv_callback,
91                    self.sudata, self.rudata, self.__exc_callback
92                )
93
94                return True
95
96            return False
97
98        def _resume(self):
99            if self.sequencer.isAlive():
100                return False
101
102            return self._start()
103       
104        def _restart(self):
105            if self.sequencer.isAlive():
106                return False
107
108            self.packet_count = 0
109            self.loop_count = 0
110            self.percentage = 0.0
111            self.answers = 0
112            self.received = 0
113
114            return self._start()
115
116        def _stop(self):
117            self.internal = False
118            self.sequencer.stop()
119
120            return True
121
122        _pause = _stop
123
124        def __exc_callback(self, exc):
125            self.internal = False
126            self.state = self.NOT_RUNNING
127            self.summary = str(exc)
128
129        def __send_callback(self, packet, want_reply, udata):
130            if not packet:
131                self.loop_count += 1
132                self.summary = _('Running sequence %d of %d times') % (self.loop_count, self.tot_loop_count)
133            else:
134                self.packet_count += 1
135
136                if want_reply:
137                    self.summary = _('Sending packet %s and waiting a reply') % packet.summary()
138                else:
139                    self.summary = _('Sending packet %s') % packet.summary()
140
141                if self.report_sent:
142                    self.data.append(packet)
143
144            self.percentage = float((float(self.packet_count) / float(self.tot_packet_count)) * \
145                                    (float(self.loop_count) / float(self.tot_loop_count))) * 100.0
146
147            if self.scallback:
148                # FIXME: THIS FUCKING UDATA also in other files
149                self.scallback(packet, want_reply, self.loop_count, self.packet_count, udata)
150
151            if not self.internal:
152                self.state = self.NOT_RUNNING
153
154            return self.state == self.NOT_RUNNING or \
155                   self.state == self.PAUSED
156
157        def __recv_callback(self, packet, reply, is_reply, udata):
158            if reply is None:
159                self.internal = False
160                self.summary = _('Sequence finished with %d packets sent and %d received') % \
161                                (self.packet_count * self.loop_count, self.received * self.loop_count)
162            else:
163                self.received += 1
164                self.summary = _('Received %s') % reply.summary()
165
166                if is_reply or self.report_recv:
167                    self.data.append(reply)
168
169            if self.rcallback:
170                self.rcallback(packet, reply, udata)
171
172            if not self.internal:
173                self.state = self.NOT_RUNNING
174
175            return self.state == self.NOT_RUNNING or \
176                   self.state == self.PAUSED
177
178        def join(self):
179            if self.sequencer.isAlive():
180                self.sequencer.stop()
181
182    return SequenceContext
Note: See TracBrowser for help on using the browser.