├── .gitignore ├── CHANGELOG ├── LICENCE ├── README.md ├── confs ├── README.md ├── argus.conf ├── example_zeo.conf ├── example_zodb_file_database.conf ├── example_zodb_server_database.conf ├── ra.conf ├── stf.conf ├── zeo.conf └── zodb.conf ├── database └── README.md ├── dependencies.txt ├── doc └── labels.md ├── modules ├── __init__.py ├── distances_1.py ├── dns_parser.py ├── example.py ├── experiments_1.py ├── markov_models_1.py ├── template_module.py └── visualize_1.py ├── stf.py └── stf ├── __init__.py ├── __init__.pyc ├── common ├── __init__.py ├── __init__.pyc ├── abstracts.py ├── ap.py ├── colors.py ├── markov_chains.py └── out.py └── core ├── __init__.py ├── __init__.pyc ├── configuration.py ├── connections.py ├── database.py ├── dataset.py ├── file.py ├── labels.py ├── models.py ├── models_constructors.py ├── notes.py ├── plugins.py └── ui ├── __init__.py ├── commands.py └── console.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | .*.swp 3 | database/stf* 4 | database/*log 5 | -------------------------------------------------------------------------------- /CHANGELOG: -------------------------------------------------------------------------------- 1 | 2 | # Stratosphere Testing Framework 0.1.3alpha (2015-4-22) 3 | - New labeling process 4 | - New modules can be loaded 5 | 6 | # Stratosphere Testing Framework 0.1.2alpha () 7 | # 8 | # Stratosphere Testing Framework 0.1alpha (2015-3-06) 9 | - Initial release 10 | 11 | # Stratosphere Testing Framework started (2015-02-06) 12 | -------------------------------------------------------------------------------- /LICENCE: -------------------------------------------------------------------------------- 1 | The Stratosphere Testing Framework is free software: you can redistribute it and/or modify 2 | it under the terms of the GNU General Public License as published by 3 | the Free Software Foundation, either version 3 of the License, or 4 | (at your option) any later version. 5 | 6 | This program is distributed in the hope that it will be useful, 7 | but WITHOUT ANY WARRANTY; without even the implied warranty of 8 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 9 | GNU General Public License for more details. 10 | 11 | You should have received a copy of the GNU General Public License 12 | along with this program. If not, see . 13 | 14 | 15 | 16 | 17 | ----------------------------------- 18 | 19 | The code taken from Viper has the following licence: 20 | Copyright (c) 2013, Claudio "nex" Guarnieri 21 | All rights reserved. 22 | 23 | Redistribution and use in source and binary forms, with or without modification, 24 | are permitted provided that the following conditions are met: 25 | 26 | * Redistributions of source code must retain the above copyright notice, this 27 | list of conditions and the following disclaimer. 28 | 29 | * Redistributions in binary form must reproduce the above copyright notice, this 30 | list of conditions and the following disclaimer in the documentation and/or 31 | other materials provided with the distribution. 32 | 33 | * Neither the name of the {organization} nor the names of its 34 | contributors may be used to endorse or promote products derived from 35 | this software without specific prior written permission. 36 | 37 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 38 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 39 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 40 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR 41 | ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 42 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 43 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON 44 | ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 45 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 46 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 47 | -------------------------------------------------------------------------------- /confs/README.md: -------------------------------------------------------------------------------- 1 | # Configuration files 2 | 3 | ## Main STF configuration 4 | The main stf configuration is stf.conf. This file contains the paths to the other configuration files needed for the databases and programs. 5 | 6 | ## Argus configuration files 7 | The configuration files for argus and ra are the following. They will be used by the execution of argus and ra in the stf program. These configuration are special to get bidirectional flows and to store some payload data inside the flow. 8 | - argus.conf 9 | - ra.conf 10 | 11 | ## ZEO database server 12 | The ZEO server is a program that listens in an ip:port and serves the object-oriented database. It is very helpful to run this program remotely. Its configuration is **zeo.conf**. The example configuration file is *example_zeo.conf*. 13 | 14 | ## ZODB configuration 15 | Since ZODB can be used with a server or directly with a file on disk, this configuration instructs zodb on how to access it. The conf file is **zodb.conf**. There are two main types that we may use. The file type of database, with an example configuration in *example_zodb_file_database.conf*, and the server type of dataset, with an example configuration in *example_zodb_server_database.conf*. 16 | -------------------------------------------------------------------------------- /confs/argus.conf: -------------------------------------------------------------------------------- 1 | # 2 | # Argus Software 3 | # Copyright (c) 2000-2011 QoSient, LLC 4 | # All rights reserved. 5 | # 6 | # Example argus.conf 7 | # 8 | # Argus will open this argus.conf if its installed as /etc/argus.conf. 9 | # It will also search for this file as argus.conf in directories 10 | # specified in $ARGUSPATH, or $ARGUSHOME, $ARGUSHOME/lib, 11 | # or $HOME, $HOME/lib, and parse it to set common configuration 12 | # options. All values in this file can be overriden by command 13 | # line options, or other files of this format that can be read in 14 | # using the -F option. 15 | # 16 | # 17 | # Variable Syntax 18 | # 19 | # Variable assignments must be of the form: 20 | # 21 | # VARIABLE= 22 | # 23 | # with no white space between the VARIABLE and the '=' sign. 24 | # Quotes are optional for string arguements, but if you want 25 | # to embed comments, then quotes are required. 26 | # 27 | # 28 | # Variable Explanations 29 | # 30 | # The Argus can be configured to support a large number of 31 | # flow types. The Argus can provide either type, i.e. 32 | # uni-directional or bi-directional flow tracking and 33 | # the flow can be further defined by specifying the key. 34 | # The argus supports a set of well known key strategies, 35 | # such as 'CLASSIC_5_TUPLE', 'LAYER_3_MATRIX', 'LAYER_2_MATRIX', 36 | # 'MPLS', and/or 'VLAN', or the argus can be configured to 37 | # formulate key strategies from a list of the specific 38 | # objects that the Argus understands. See the man page for 39 | # a complete description. 40 | # 41 | # The default is the classic 5-tuple IP flow, CLASSIC_5_TUPLE. 42 | # 43 | 44 | ARGUS_FLOW_TYPE="Bidirectional" 45 | # For botnet-cluster is NOT commented ARGUS_FLOW_TYPE="Unidirectional" 46 | #ARGUS_FLOW_TYPE="Unidirectional" 47 | ARGUS_FLOW_KEY="CLASSIC_5_TUPLE" 48 | 49 | 50 | # Argus is capable of running as a daemon, doing all the right things 51 | # that daemons do. When this configuration is used for the system 52 | # daemon process, say for /etc/argus.conf, this variable should be 53 | # set to "yes". 54 | # 55 | # The default value is to not run as a daemon. 56 | # 57 | # This example is to support the ./support/Startup/argus script 58 | # which works when this variable be set to "yes". 59 | # 60 | # Commandline equivalent -d 61 | # 62 | 63 | #ARGUS_DAEMON=no 64 | 65 | 66 | # Argus Monitor Data is uniquely identifiable based on the source 67 | # identifier that is included in each output record. This is to 68 | # allow you to work with Argus Data from multiple monitors at the 69 | # same time. The ID is 32 bits long, and so legitimate values are 70 | # 0 - 4294967296 but argus also supports IP addresses as values. 71 | # The configuration allows for you to use host names, however, do 72 | # have some understanding how `hostname` will be resolved by the 73 | # nameserver before commiting to this strategy completely. 74 | # 75 | # For convenience, argus supports the notion of "`hostname`" for 76 | # assigning the probe's id. This is to support management of 77 | # large deployments, so you can have one argus.conf file that works 78 | # for a lot of probes. 79 | # 80 | # For security, argus does not rely on system programs, like hostname.1, 81 | # It implements the logic of hostname itself, so don't try to run 82 | # arbitrary programs using this method, because it won't work. 83 | # 84 | # Commandline equivalent -e 85 | # 86 | 87 | ARGUS_MONITOR_ID=`hostname` 88 | 89 | 90 | # Argus monitors can provide a real-time remote access port 91 | # for collecting Argus data. This is a TCP based port service and 92 | # the default port number is tcp/561, the "experimental monitor" 93 | # service. This feature is disabled by default, and can be forced 94 | # off by setting it to zero (0). 95 | # 96 | # When you do want to enable this service, 561 is a good choice, 97 | # as all ra* clients are configured to try this port by default. 98 | # 99 | # Commandline equivalent -P 100 | # 101 | 102 | #ARGUS_ACCESS_PORT=561 103 | ARGUS_ACCESS_PORT=902 104 | 105 | 106 | # When remote access is enabled (see above), you can specify that Argus 107 | # should bind only to a specific IP address. This is useful, for example, 108 | # in restricting access to the local host, or binding to a private 109 | # interface while capturing from another. 110 | # 111 | # You can provide multiple addresses, separated by commas, or on multiple 112 | # lines. 113 | # 114 | # The default is to bind to any IP address. 115 | # 116 | # Commandline equivalent -B 117 | # 118 | 119 | #ARGUS_BIND_IP="::1,127.0.0.1" 120 | #ARGUS_BIND_IP="127.0.0.1" 121 | #ARGUS_BIND_IP="192.168.0.68" 122 | 123 | 124 | # By default, Argus will open the first appropriate interface on a 125 | # system that it encounters. For systems that have only one network 126 | # interface, this is a reasonable thing to do. But, when there are 127 | # more than one suitable interface, you should specify the 128 | # interface(s) Argus should use either on the command line or in this 129 | # file. 130 | # 131 | # Argus can track packets from any or all interfaces, concurrently. 132 | # The interfaces can be tracked as: 133 | # 1. independant - this is where argus tracks flows from each 134 | # interface independant from the packets seen on any other 135 | # interface. This is useful for hosts/routers that 136 | # have full-duplex interfaces, and you want to distinguish 137 | # flows based on their interface. There is an option to specify 138 | # a distinct srcid to each independant modeler. 139 | # 140 | # 2. duplex - where argus tracks packets from 2 interfaces 141 | # as if they were two half duplex streams of the same link. 142 | # Because there is a single modeler tracking the 2 143 | # interfaces, there is a single srcid that can be passed as 144 | # an option. 145 | # 146 | # 3. bonded - where argus tracks packets from multiple interfaces 147 | # as if they were from the same stream. Because there is a 148 | # single modeler tracking the 2 interfaces, there is a single 149 | # srcid that can be passed as an option. 150 | # 151 | # Interfaces can be specified as groups using '[',']' notation, to build 152 | # flexible definitions of packet sources. However, each interface 153 | # should be referenced only once (this is due to performance and OS 154 | # limitations, so if your OS has no problem with this, go ahead). 155 | # 156 | # The lo (loopback) interface will be included only if it is specifically 157 | # indicated in the option. 158 | # 159 | # The syntax for specifying this either on the command line or in this file: 160 | # -i ind:all 161 | # -i dup:en0,en1/srcid 162 | # -i bond:en0,en1/srcid 163 | # -i dup:[bond:en0,en1],en2/srcid 164 | # -i en0/srcid -i en1/srcid (equivalent '-i ind:en0/srcid,en1/srcid') 165 | # -i en0 en1 (equivalent '-i bond:en0,en1') 166 | # -i en1(dlt)/srcid -i en1(dlt)/srcid 167 | # 168 | # In all cases, if there is a "-e srcid" provided, the srcid provided is used 169 | # as the default. If a srcid is specified using this option, it overrides 170 | # the default. 171 | # 172 | # Commandline equivalent -i 173 | # 174 | 175 | #ARGUS_INTERFACE=any 176 | #ARGUS_INTERFACE=ind:all 177 | #ARGUS_INTERFACE=ind:en0/192.168.0.68,en2/192.168.2.1 178 | #ARGUS_INTERFACE=en0 179 | 180 | 181 | # By default, Argus will put its interface in promiscuous mode 182 | # in order to monitor all the traffic that can be collected. 183 | # This can put an undo load on systems. 184 | 185 | # If the intent is to monitor only the network activity of 186 | # the specific system, say to measure the performance of 187 | # an HTTP service or DNS service, you'll want to turn 188 | # promiscuous mode off. 189 | # 190 | # The default value is go into prmiscuous mode. 191 | # 192 | # Commandline equivalent -p 193 | # 194 | 195 | #ARGUS_GO_PROMISCUOUS=yes 196 | 197 | 198 | # Argus supports chroot(2) in order to control the file system that 199 | # argus exists in and can access. Generally used when argus is running 200 | # with privileges, this limits the negative impacts that argus could 201 | # inflict on its host machine. 202 | # 203 | # This option will cause the output file names to be relative to this 204 | # directory, and so consider this when trying to find your output files. 205 | # 206 | # Commandline equivalent -c 207 | # 208 | 209 | #ARGUS_CHROOT_DIR=/chroot_dir 210 | 211 | 212 | # Argus can be configured to enable detailed control plane 213 | # flow monitoring for specific control plane protocols. 214 | # 215 | # This feature requires full packet capture for the monitored 216 | # interface in order to capture the complete control plane 217 | # protocol, and will have a performance impact on the sensor. 218 | # 219 | # The default is to not turn this feature on. 220 | # 221 | # Commandline equivalent -C 222 | # 223 | 224 | #ARGUS_CAPTURE_FULL_CONTROL_DATA=no 225 | # For botnet-cluster is commented ARGUS_CAPTURE_FULL_CONTROL_DATA=no 226 | ARGUS_CAPTURE_FULL_CONTROL_DATA=yes 227 | 228 | 229 | # Argus can be directed to change its user id using the setuid() system 230 | # call. This is can used when argus is started as root, in order to 231 | # access privileged resources, but then after the resources are opened, 232 | # this directive will cause argus to change its user id value to 233 | # a 'lesser' capable account. Recommended when argus is running as 234 | # daemon. 235 | # 236 | # Commandline equivalent -u 237 | # 238 | 239 | #ARGUS_SETUSER_ID=user 240 | 241 | 242 | # Argus can be directed to change its group id using the setgid() system 243 | # call. This is can used when argus is started as root, in order to 244 | # access privileged resources, but then after the resources are opened, 245 | # this directive can be used to change argu's group id value to 246 | # a 'lesser' capable account. Recommended when argus is running as 247 | # daemon. 248 | # 249 | # Commandline equivalent -g 250 | # 251 | 252 | #ARGUS_SETGROUP_ID=group 253 | 254 | 255 | # Argus can write its output to one or a number of files. 256 | # The default limit is 5 concurrent files, each with their 257 | # own independant filters. 258 | # 259 | # The format is: 260 | # ARGUS_OUTPUT_FILE=/full/path/file/name 261 | # ARGUS_OUTPUT_FILE="/full/path/file/name filter" 262 | # 263 | # Most sites will have argus write to a file, for reliablity. 264 | # The example file name is used here as supporting programs, 265 | # such as ./support/Archive/argusarchive are configured to use 266 | # this file (with any chroot'd directory prepended). 267 | # 268 | # Commandline equivalent -w 269 | # 270 | 271 | #ARGUS_OUTPUT_FILE=/var/log/argus/argus.out 272 | 273 | 274 | # Argus can push its output to one or a number of remote hosts. 275 | # The default limit is 5 concurrent output streams, each with their 276 | # own independant filters. 277 | # 278 | # The format is: 279 | # ARGUS_OUTPUT_STREAM="URI [filter]" 280 | # ARGUS_OUTPUT_STREAN="argus-udp://multicastGroup:port 281 | # ARGUS_OUTPUT_STREAN="argus-udp://host:port 'tcp and not udp'" 282 | # 283 | # Most sites will have argus listen() for remote sites to request argus data, 284 | # using a "pull" data model. But for some sites and applications, pushing 285 | # records without explicit registration is desired. This option will cause 286 | # argus to transmit records that match the optional filter, to the configured 287 | # targets using UDP as the transport mechanism. 288 | # 289 | # The primary purpose for this feature is to multicast argus records to 290 | # a number of listeners on an interface, but it is not limited to this 291 | # purpose. The multicast TTL is set to 128 by default, so that you can 292 | # send records some distance. 293 | # 294 | # Commandline equivalent -w argus-udp://host:port 295 | # 296 | 297 | #ARGUS_OUTPUT_STREAM=argus-udp://224.0.20.21:561 298 | 299 | 300 | # When Argus is configured to run as a daemon, with the -d 301 | # option, Argus can store its pid in a file, to aid in 302 | # managing the running daemon. However, creating a system 303 | # pid file requires priviledges that may not be appropriate 304 | # for all cases. 305 | # 306 | # When configured to generate a pid file, if Argus cannot 307 | # create the pid file, it will fail to run. This variable 308 | # is available to override the default, in case this gets 309 | # in your way. 310 | # 311 | # The default value is to generate a pid. The default 312 | # path for the pid file, is '/var/run'. 313 | # 314 | # No Commandline equivalent 315 | # 316 | 317 | ARGUS_SET_PID=yes 318 | ARGUS_PID_PATH="/var/run" 319 | 320 | 321 | # Argus will periodically report on a flow's activity every 322 | # ARGUS_FLOW_STATUS_INTERVAL seconds, as long as there is 323 | # new activity on the flow. This is so that you can get a 324 | # multiple status reports into the activity of a flow. The 325 | # default is 5 seconds, but this number may be too low or 326 | # too high depending on your uses. Argus does suppport 327 | # a minimum value of 0.000001 seconds. Values under 1 sec 328 | # are very useful for doing measurements in a controlled 329 | # experimental environment where the number of flows is small. 330 | # 331 | # Because the status interval affects the memory utilization 332 | # of the monitor, find the minimum acceptable value is 333 | # recommended. 334 | # 335 | # Commandline equivalent -S 336 | # 337 | 338 | #ARGUS_FLOW_STATUS_INTERVAL=5 339 | ARGUS_FLOW_STATUS_INTERVAL=3600 340 | 341 | 342 | # Argus will periodically report on a its own health, providing 343 | # interface status, total packet and bytes counts, packet drop 344 | # rates, and flow oriented statistics. 345 | # 346 | # These records can be used as "keep alives" for periods when 347 | # there is no network traffic to be monitored. 348 | # 349 | # The default value is 300 seconds, but a value of 60 seconds is 350 | # very common. 351 | # 352 | # Commandline equivalent -M 353 | # 354 | 355 | ARGUS_MAR_STATUS_INTERVAL=60 356 | 357 | 358 | # Argus has a number of flow state timers that specify how long argus 359 | # will 'remember' the caches of specific flows after they have gone 360 | # idle. 361 | # 362 | # The default values have been chosen to aggresively timeout flow 363 | # caches to conserve memory utilization. Increasing values can have 364 | # an impact on argus memory use, so take care when modifying values. 365 | # 366 | # If you think there is a flow type that doesn't have appropriate 367 | # timeout support, send email to the developer's list, we'll add one 368 | # for you. 369 | # 370 | 371 | #ARGUS_IP_TIMEOUT=30 372 | #ARGUS_TCP_TIMEOUT=60 373 | #ARGUS_ICMP_TIMEOUT=5 374 | #ARGUS_IGMP_TIMEOUT=30 375 | #ARGUS_FRAG_TIMEOUT=5 376 | #ARGUS_ARP_TIMEOUT=5 377 | #ARGUS_OTHER_TIMEOUT=30 378 | 379 | 380 | # If compiled to support this option, Argus is capable of 381 | # generating a lot of debug information. 382 | # 383 | # The default value is zero (0). 384 | # 385 | # Commandline equivalent -D 386 | # 387 | 388 | ARGUS_DEBUG_LEVEL=0 389 | 390 | 391 | # Argus can be configured to report on flows in a manner than 392 | # provides the best information for calculating application 393 | # reponse times and network round trip times. 394 | # 395 | # The default value is to not generate this data. 396 | # 397 | # Commandline equivalent -R 398 | # 399 | 400 | # for botnet-clusterer it was no ARGUS_GENERATE_RESPONSE_TIME_DATA=no 401 | #ARGUS_GENERATE_RESPONSE_TIME_DATA=yes 402 | ARGUS_GENERATE_RESPONSE_TIME_DATA=yes 403 | 404 | 405 | # Argus can be configured to generate packet size information 406 | # on a per flow basis, which provides the max and min packet 407 | # size seen . The default value is to not generate this data. 408 | # 409 | # Commandline equivalent -Z 410 | # 411 | 412 | ARGUS_GENERATE_PACKET_SIZE=yes 413 | 414 | 415 | # Argus can be configured to generate packet jitter information 416 | # on a per flow basis. The default value is to not generate 417 | # this data. 418 | # 419 | # Commandline equivalent -J 420 | # 421 | 422 | # for botnet-cluster it was no ARGUS_GENERATE_JITTER_DATA=no 423 | #ARGUS_GENERATE_JITTER_DATA=yes 424 | ARGUS_GENERATE_JITTER_DATA=yes 425 | 426 | 427 | # Argus can be configured to provide MAC addresses in 428 | # it audit data. The default value is to not generate 429 | # this data. 430 | # 431 | # Commandline equivalent -m 432 | # 433 | 434 | ARGUS_GENERATE_MAC_DATA=yes 435 | 436 | 437 | # Argus can be configured to generate metrics that include 438 | # the application byte counts as well as the packet count 439 | # and byte counters. 440 | # 441 | # Commandline equivalent -A 442 | # 443 | 444 | ARGUS_GENERATE_APPBYTE_METRIC=yes 445 | 446 | 447 | # Argus by default, generates extended metrics for TCP 448 | # that include the connection setup time, window sizes, 449 | # base sequence numbers, and retransmission counters. 450 | # You can suppress this detailed information using this 451 | # variable. 452 | # 453 | # No commandline equivalent 454 | # 455 | 456 | # for botnet-clusterer was commented ARGUS_GENERATE_TCP_PERF_METRIC=yes 457 | #ARGUS_GENERATE_TCP_PERF_METRIC=yes 458 | ARGUS_GENERATE_TCP_PERF_METRIC=yes 459 | 460 | 461 | # Argus by default, generates a single pair of timestamps, 462 | # for the first and last packet seen on a given flow, during 463 | # the obseration period. For bi-directional flows, this 464 | # results in loss of some information. By setting this 465 | # variable to 'yes', argus will store start and ending 466 | # timestamps for both directions of the flow. 467 | # 468 | # No commandline equivalent 469 | # 470 | 471 | ARGUS_GENERATE_BIDIRECTIONAL_TIMESTAMPS=yes 472 | 473 | 474 | # Argus can be configured to capture a number of user data 475 | # bytes from the packet stream. 476 | # 477 | # The default value is to not generate this data. 478 | # 479 | # Commandline equivalent -U 480 | # 481 | 482 | ARGUS_CAPTURE_DATA_LEN=480 483 | 484 | 485 | 486 | 487 | # Argus uses the packet filter capabilities of libpcap. If 488 | # there is a need to not use the libpcap filter optimizer, 489 | # you can turn it off here. The default is to leave it on. 490 | # 491 | # Commandline equivalent -O 492 | # 493 | 494 | #ARGUS_FILTER_OPTIMIZER=yes 495 | 496 | 497 | # You can provide a filter expression here, if you like. 498 | # It should be limited to 2K in length. The default is to 499 | # not filter. 500 | # 501 | # The commandline filter will override this filter expression. 502 | # 503 | 504 | #ARGUS_FILTER="" 505 | 506 | 507 | # Argus allows you to capture packets in tcpdump() format 508 | # if the source of the packets is a tcpdump() formatted 509 | # file or live packet source. 510 | # 511 | # Specify the path to the packet capture file here. 512 | # 513 | 514 | #ARGUS_PACKET_CAPTURE_FILE="/var/log/argus/packet.out" 515 | 516 | 517 | # Argus supports the use of SASL to provide strong 518 | # authentication and confidentiality protection. 519 | # 520 | # The policy that argus uses is controlled through 521 | # the use of a minimum and maximum allowable protection 522 | # strength. Set these variable to control this policy. 523 | # 524 | 525 | #ARGUS_MIN_SSF=40 526 | #ARGUS_MAX_SSF=128 527 | 528 | 529 | # Argus supports setting environment variables to enable 530 | # functions required by the kernel or shared libraries. 531 | # This feature is intended to support libraries such as 532 | # the net pf_ring support for libpcap as supported by 533 | # code at http://public.lanl.gov/cpw/ 534 | # 535 | # Setting environment variables in this way does not affect 536 | # internal argus variable in any way. As a result, you 537 | # can't set ARGUS_PATH using this feature. 538 | # 539 | # Care should must be taken to assure that the value given 540 | # the variable conform's to your systems putenv.3 system call. 541 | # You can have as many of these directives as you like. 542 | # 543 | # The example below is intended to set a libpcap ring buffer 544 | # length to 300MB, if your system supports this feature. 545 | # 546 | 547 | #ARGUS_ENV="PCAP_MEMORY=300000" 548 | 549 | 550 | # Argus can be configured to discover tunneling protocols 551 | # above the UDP transport header, specifically Teredo 552 | # (IPv6 over UDP). The algorithm is simple and so, having 553 | # this on by default may generate false tunnel matching. 554 | 555 | # The default is to not turn this feature on. 556 | 557 | #ARGUS_TUNNEL_DISCOVERY="no" 558 | 559 | 560 | 561 | # Argus can be configured to be self synchronizing with other 562 | # argi. This involves using state from packets contents to 563 | # synchronize the flow reporting. 564 | # 565 | 566 | #ARGUS_SELF_SYNCHRONIZE=yes 567 | 568 | 569 | 570 | # Argus supports the generation of host originated processes 571 | # to gather additional data and statistics. These include 572 | # periodic processes to poll for SNMP data, as an example, or 573 | # to collect host statistics through reading procfs(). Or 574 | # single run programs that run at a specified time. 575 | # 576 | # These argus events, are generated from the complete list of 577 | # ARGUS_EVENT_DATA directives that are specified here. 578 | # 579 | # The syntax is: 580 | # Syntax is: "method:path|prog:interval[:postproc]" 581 | # Where: method = [ "file" | "prog" ] 582 | # pathname | program = "%s" 583 | # interval = %d[smhd] [ zero means run once ] 584 | # postproc = [ "compress" | "compress2" ] 585 | # 586 | #ARGUS_EVENT_DATA="prog:/usr/local/bin/ravms:20s:compress" 587 | #ARGUS_EVENT_DATA="prog:/usr/local/bin/rasnmp:1m:compress" 588 | #ARGUS_EVENT_DATA="file:/proc/vmstat:30s:compress" 589 | #ARGUS_EVENT_DATA="prog:/usr/bin/uptime:30s" 590 | #ARGUS_EVENT_DATA="prog:/usr/local/bin/ralsof:30s:compress" 591 | 592 | 593 | # This version of Argus supports keystroke detection and counting for 594 | # TCP connections, with specific algoritmic support for SSH connections. 595 | # 596 | # The ARGUS_KEYSTROKE variable turns the feature on. Values for 597 | # this variable are: 598 | # ARGUS_KEYSTROKE="yes" - turn on TCP flow tracking 599 | # ARGUS_KEYSTROKE="tcp" - turn on TCP flow tracking 600 | # ARGUS_KEYSTROKE="ssh" - turn on SSH specific flow tracking 601 | # ARGUS_KEYSTROKE="no" [default] 602 | # 603 | # The algorithm uses a number of variables, all of which can be 604 | # modifed using the ARGUS_KEYSTROKE_CONF descriptor, which is a 605 | # semicolon (';') separated set of variable assignments. Here is 606 | # the list of supported variables: 607 | # 608 | # DC_MIN - (int) Minimum client datagram payload size in bytes 609 | # DC_MAX - (int) Maximum client datagram payload size in bytes 610 | # GS_MAX - (int) Maximum server packet gap 611 | # DS_MIN - (int) Minimum server datagram payload size in bytes 612 | # DS_MAX - (int) Maximum server datagram payload size in bytes 613 | # IC_MIN - (int) Minimum client interpacket arrival time (microseconds) 614 | # LCS_MAX - (int) Maximum something - Not sure what this is 615 | # GPC_MAX - (int) Maximum client packet gap 616 | # ICR_MIN - (float) Minimum client/server interpacket arrival ratio 617 | # ICR_MAX - (float) Maximum client/server interpacket arrival ratio 618 | # 619 | # All variables have default values, this variable is used to override 620 | # those values. The syntax for the variable is: 621 | # 622 | # botnet-clusterer was commentd ARGUS_KEYSTROKE_CONF="DC_MIN=20;DS_MIN=20" 623 | ARGUS_KEYSTROKE_CONF="DC_MIN=20;DS_MIN=20" 624 | # 625 | 626 | #botnet-clusterer was no ARGUS_KEYSTROKE="no" 627 | ARGUS_KEYSTROKE="ssh" 628 | #ARGUS_KEYSTROKE="no" 629 | #ARGUS_KEYSTROKE_CONF="" 630 | -------------------------------------------------------------------------------- /confs/example_zeo.conf: -------------------------------------------------------------------------------- 1 | 2 | address localhost:9002 3 | 4 | 5 | 6 | path ./database/stf.db 7 | 8 | 9 | 10 | 11 | path ./database/zeo.log 12 | format %(asctime)s %(message)s 13 | 14 | 15 | 16 | -------------------------------------------------------------------------------- /confs/example_zodb_file_database.conf: -------------------------------------------------------------------------------- 1 | 2 | 3 | path database/stf.db 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /confs/example_zodb_server_database.conf: -------------------------------------------------------------------------------- 1 | 2 | 3 | server localhost:9002 4 | 5 | 6 | -------------------------------------------------------------------------------- /confs/ra.conf: -------------------------------------------------------------------------------- 1 | # 2 | # Argus Software 3 | # Copyright (c) 2000-2011 QoSient, LLC 4 | # All rights reserved. 5 | # 6 | # This program is free software; you can redistribute it and/or modify 7 | # it under the terms of the GNU General Public License as published by 8 | # the Free Software Foundation; either version 2, or (at your option) 9 | # any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU General Public License 17 | # along with this program; if not, write to the Free Software 18 | # Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA. 19 | # 20 | # 21 | # Ra.Print.All.Conf 22 | # 23 | # This ra rc file will print all the available fields for a given argus 24 | # record, seperated by a comma. 25 | #RA_FIELD_SPECIFIER= srcid seq stime ltime dur sstime sltime sdur dstime dltime ddur srng drng trans flgs avgdur stddev mindur maxdur saddr dir daddr proto sport dport sco dco stos dtos sdsb ddsb sttl dttl shops dhops sipid dipid pkts spkts dpkts bytes sbytes dbytes appbytes sappbytes dappbytes load sload dload rate srate drate loss sloss dloss ploss sploss dploss senc denc smac dmac smpls dmpls svlan dvlan svid dvid svpri dvpri sintpkt dintpkt sintpktact dintpktact sintpktidl dintpktidl sintpktmax sintpktmin dintpktmax dintpktmin sintpktactmax sintpktactmin dintpktactmax dintpktactmin sintpktidlmax sintpktidlmin dintpktidlmax dintpktidlmin jit sjit djit jitact sjitact djitact jitidl sjitidl djitidl state deldur delstime delltime dspkts ddpkts dsbytes ddbytes pdspkts pddpkts pdsbytes pddbytes suser:1500 duser:1500 tcpext swin dwin jdelay ldelay bins binnum stcpb dtcpb tcprtt synack ackdat inode smaxsz sminsz dmaxsz dminsz 26 | RA_PRINT_LABELS=0 27 | RA_USEC_PRECISION=6 28 | RA_PRINT_NAMES=0 29 | RA_TIME_FORMAT="%Y/%m/%d %T.%f" 30 | 31 | RA_FIELD_DELIMITER=',' 32 | RA_USERDATA_ENCODE=Encode64 33 | # stf without the fields of suser and duser 34 | #RA_FIELD_SPECIFIER= stime dur proto:10 saddr:27 sport dir daddr:27 dport state stos dtos pkts bytes sbytes label:40 35 | # stf conf 36 | RA_FIELD_SPECIFIER= stime dur proto:10 saddr:27 sport dir daddr:27 dport state stos dtos pkts bytes sbytes suser:480 duser:480 label:40 37 | 38 | -------------------------------------------------------------------------------- /confs/stf.conf: -------------------------------------------------------------------------------- 1 | # The stf section 2 | [stf] 3 | # Configuration file for the zeo server database, in case it is used 4 | zeoconfigurationfile = confs/zeo.conf 5 | # Configuration file for the zodb database. Normally this is the only one used 6 | zodbconfigurationfile = confs/zodb.conf 7 | -------------------------------------------------------------------------------- /confs/zeo.conf: -------------------------------------------------------------------------------- 1 | # Usually you don't need to touch this file unless you are using runzeo to run a database server 2 | 3 | address localhost:9002 4 | 5 | 6 | 7 | path ./database/stf.db 8 | 9 | 10 | 11 | 12 | path ./database/zeo.log 13 | format %(asctime)s %(message)s 14 | 15 | 16 | 17 | -------------------------------------------------------------------------------- /confs/zodb.conf: -------------------------------------------------------------------------------- 1 | 2 | 3 | path database/stf.db 4 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /database/README.md: -------------------------------------------------------------------------------- 1 | # Database folder 2 | In this folder the program will store the database when is used. Please do not delete it. 3 | -------------------------------------------------------------------------------- /dependencies.txt: -------------------------------------------------------------------------------- 1 | python-prettytable 2 | python-transaction 3 | python-persistent 4 | python-zodb 5 | python-sparse 6 | 7 | -------------------------------------------------------------------------------- /doc/labels.md: -------------------------------------------------------------------------------- 1 | # How to work with labels in stf 2 | 3 | The assignment of labels in stf is simple. However, it is not simple to decide which label corresponds to a connection. The problem cames from the definition of label to us. We want our labels to be very precise and descriptive. That is why the ```label``` command in stf is asking for some questions before putting a label. The questions are: 4 | 5 | - Which is the direction of the connection, respect to the main host on it? 6 | This question is directed to know if the data is comming _From_ or _To_ a _Normal_ or _Botnet_ host. The directionality is important to create different behavioral models later. For example, not all the data comming _to_ an infected computer should be labeled as _Botnet_. 7 | 8 | - Which is the main decision on the data? 9 | This question is basically to know if you consider the data as _Botnet_, _Normal_, etc. 10 | 11 | - Which is the main protocol up to layer 4? 12 | In this question we want to know which is in your opinion the main protocol that is characteristic of this connection. It can be HTTP, but it can also be P2P or ARP, depending on the connection. 13 | 14 | - A description. 15 | This description is useful to separate traffic from different websites for example, or different malware families. 16 | 17 | We this information we _may_ be able to get a label name, however we have another issue to consider. It is common that two connections have the same origin, decision, main protocol and description, like for example **From-Botnet-UDP-DNS-DGA**. However, two DGA connections can have very different behavioral patterns. To distinguish them we add a number to the end of the label, in order to have **From-Botnet-UDP-DNS-DGA** and **From-Botnet-UDP-DNS-DGA**. 18 | -------------------------------------------------------------------------------- /modules/__init__.py: -------------------------------------------------------------------------------- 1 | # This file is part of Viper - https://github.com/botherder/viper 2 | # See the file 'LICENSE' for copying permission. 3 | -------------------------------------------------------------------------------- /modules/dns_parser.py: -------------------------------------------------------------------------------- 1 | # Part of this file was taken from Viper - https://github.com/botherder/viper 2 | # The rest is from the Stratosphere Testing Framework 3 | # See the file 'LICENSE' for copying permission. 4 | 5 | # Example file of how to create a module without persistence in the database. Useful for obtaining statistics or processing data. 6 | 7 | from stf.common.out import * 8 | from stf.common.abstracts import Module 9 | from stf.core.models import __groupofgroupofmodels__ 10 | from stf.core.dataset import __datasets__ 11 | from stf.core.notes import __notes__ 12 | from stf.core.connections import __group_of_group_of_connections__ 13 | from stf.core.models_constructors import __modelsconstructors__ 14 | from stf.core.labels import __group_of_labels__ 15 | import tempfile 16 | from subprocess import Popen 17 | from subprocess import PIPE 18 | 19 | 20 | import pprint 21 | import socket 22 | import struct 23 | 24 | 25 | def decode_labels(message, offset): 26 | labels = [] 27 | 28 | while True: 29 | length, = struct.unpack_from("!B", message, offset) 30 | 31 | if (length & 0xC0) == 0xC0: 32 | pointer, = struct.unpack_from("!H", message, offset) 33 | offset += 2 34 | 35 | return labels + decode_labels(message, pointer & 0x3FFF), offset 36 | 37 | if (length & 0xC0) != 0x00: 38 | raise StandardError("unknown label encoding") 39 | 40 | offset += 1 41 | 42 | if length == 0: 43 | return labels, offset 44 | 45 | labels.append(*struct.unpack_from("!%ds" % length, message, offset)) 46 | offset += length 47 | 48 | 49 | DNS_QUERY_SECTION_FORMAT = struct.Struct("!2H") 50 | 51 | def decode_question_section(message, offset, qdcount): 52 | questions = [] 53 | 54 | for _ in range(qdcount): 55 | qname, offset = decode_labels(message, offset) 56 | 57 | qtype, qclass = DNS_QUERY_SECTION_FORMAT.unpack_from(message, offset) 58 | offset += DNS_QUERY_SECTION_FORMAT.size 59 | 60 | question = {"domain_name": qname, 61 | "query_type": qtype, 62 | "query_class": qclass} 63 | 64 | questions.append(question) 65 | 66 | return questions, offset 67 | 68 | 69 | DNS_QUERY_MESSAGE_HEADER = struct.Struct("!6H") 70 | 71 | def decode_dns_message(message): 72 | 73 | id, misc, qdcount, ancount, nscount, arcount = DNS_QUERY_MESSAGE_HEADER.unpack_from(message) 74 | 75 | qr = (misc & 0x8000) != 0 76 | opcode = (misc & 0x7800) >> 11 77 | aa = (misc & 0x0400) != 0 78 | tc = (misc & 0x200) != 0 79 | rd = (misc & 0x100) != 0 80 | ra = (misc & 0x80) != 0 81 | z = (misc & 0x70) >> 4 82 | rcode = misc & 0xF 83 | 84 | offset = DNS_QUERY_MESSAGE_HEADER.size 85 | questions, offset = decode_question_section(message, offset, qdcount) 86 | 87 | result = {"id": id, 88 | "is_response": qr, 89 | "opcode": opcode, 90 | "is_authoritative": aa, 91 | "is_truncated": tc, 92 | "recursion_desired": rd, 93 | "recursion_available": ra, 94 | "reserved": z, 95 | "response_code": rcode, 96 | "question_count": qdcount, 97 | "answer_count": ancount, 98 | "authority_count": nscount, 99 | "additional_count": arcount, 100 | "questions": questions} 101 | 102 | return result 103 | 104 | 105 | #s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) 106 | #host = '' 107 | #port = 53 108 | #size = 512 109 | #s.bind((host, port)) 110 | #while True: 111 | # data, addr = s.recvfrom(size) 112 | # pprint.pprint(decode_dns_message(data)) 113 | 114 | 115 | 116 | 117 | 118 | 119 | 120 | 121 | 122 | 123 | 124 | class DNSInfo(Module): 125 | cmd = 'dns_parser' 126 | description = 'Module to compute statistics of a DNS tuple.' 127 | authors = ['Harpo MAxx'] 128 | 129 | def __init__(self): 130 | # Call to our super init 131 | super(DNSInfo, self).__init__() 132 | self.parser.add_argument('-i', '--info', metavar='tuple_id', help='Show DNS statistics for this 4-tuple.') 133 | 134 | def show_flows(self,group_of_connections,connection_id): 135 | all_text='' 136 | # Access each flow in this 4-tuple 137 | for flow_id in group_of_connections.connections[connection_id].flows: 138 | # all_text = all_text + group_of_connections.connections[connection_id].flows[flow_id].get_srcUdata() + '\n' 139 | dns_text = ":".join("{:02x}".format(ord(c)) for c in group_of_connections.connections[connection_id].flows[flow_id].get_srcUdata()) 140 | all_text = all_text + dns_text + '\n\n' 141 | f = tempfile.NamedTemporaryFile() 142 | f.write(all_text) 143 | f.flush() 144 | p = Popen('less -R ' + f.name, shell=True, stdin=PIPE) 145 | p.communicate() 146 | sys.stdout = sys.__stdout__ 147 | f.close() 148 | 149 | def dns_info(self,connection_id): 150 | # """ Show the flows inside a connection """ 151 | # #__group_of_group_of_connections__.show_flows_in_connnection(arg, filter) 152 | if __datasets__.current: 153 | group_id = __datasets__.current.get_id() 154 | try: 155 | print_info('Showing the DNS information of 4tuple {}'.format(connection_id)) 156 | self.show_flows(__group_of_group_of_connections__.group_of_connections[group_id],connection_id) 157 | except KeyError: 158 | print_error('The connection {} does not longer exists in the database.'.format(connection_id)) 159 | else: 160 | print_error('There is no dataset selected.') 161 | 162 | def run(self): 163 | # List general info 164 | def help(): 165 | self.log('info', self.description) 166 | # Run? 167 | super(DNSInfo, self).run() 168 | if self.args is None: 169 | return 170 | if self.args.info: 171 | self.dns_info(self.args.info) 172 | else: 173 | print_error('At least one parameter is required') 174 | self.usage() 175 | -------------------------------------------------------------------------------- /modules/example.py: -------------------------------------------------------------------------------- 1 | # Part of this file was taken from Viper - https://github.com/botherder/viper 2 | # The rest is from the Stratosphere Testing Framework 3 | # See the file 'LICENSE' for copying permission. 4 | 5 | # Example file of how to create a module without persistence in the database. Useful for obtaining statistics or processing data. 6 | 7 | from stf.common.out import * 8 | from stf.common.abstracts import Module 9 | from stf.core.models import __groupofgroupofmodels__ 10 | from stf.core.dataset import __datasets__ 11 | from stf.core.notes import __notes__ 12 | from stf.core.connections import __group_of_group_of_connections__ 13 | from stf.core.models_constructors import __modelsconstructors__ 14 | from stf.core.labels import __group_of_labels__ 15 | 16 | 17 | class Example(Module): 18 | cmd = 'example' 19 | description = 'Example module to print some statistics about the data in stf' 20 | authors = ['Sebastian Garcia'] 21 | 22 | def __init__(self): 23 | # Call to our super init 24 | super(Example, self).__init__() 25 | self.parser.add_argument('-i', '--info', action='store_true', help='Show info') 26 | 27 | 28 | def example_info(self): 29 | # Example to read datasets 30 | datasets_ids = list(__datasets__.get_datasets_ids()) 31 | print_info('There are {} datasets: {}.'.format(len(datasets_ids), datasets_ids)) 32 | 33 | # Example to get connnections 34 | connections_groups_ids = list(__group_of_group_of_connections__.get_groups_ids()) 35 | print_info('There are {} groups of connections: {}.'.format(len(connections_groups_ids), connections_groups_ids)) 36 | 37 | # Example to read models 38 | models_groups_ids = list(__groupofgroupofmodels__.get_groups_ids()) 39 | print_info('There are {} groups of models: {}.'.format(len(models_groups_ids), models_groups_ids)) 40 | 41 | # Example to get notes 42 | notes_ids = list(__notes__.get_notes_ids()) 43 | print_info('There are {} notes: {}.'.format(len(notes_ids), notes_ids)) 44 | 45 | # Example to get labels 46 | labels_ids = list(__group_of_labels__.get_labels_ids()) 47 | print_info('There are {} labels: {}.'.format(len(labels_ids), labels_ids)) 48 | 49 | # Example to get the labels constructors 50 | constructors_ids = list(__modelsconstructors__.get_constructors_ids()) 51 | print_info('There are {} model constructors: {}.'.format(len(constructors_ids), constructors_ids)) 52 | 53 | 54 | 55 | 56 | def run(self): 57 | # List general info 58 | def help(): 59 | self.log('info', "Example module") 60 | print 'Hi' 61 | 62 | # Run? 63 | super(Example, self).run() 64 | if self.args is None: 65 | return 66 | 67 | if self.args.info: 68 | self.example_info() 69 | else: 70 | print_error('At least one of the parameter is required') 71 | self.usage() 72 | -------------------------------------------------------------------------------- /modules/template_module.py: -------------------------------------------------------------------------------- 1 | # Part of this file was taken from Viper - https://github.com/botherder/viper 2 | # The rest is from the Stratosphere Testing Framework 3 | # See the file 'LICENSE' for copying permission. 4 | 5 | # This is a template module showing how to create a module that has persistence in the database. To create your own command just copy this file and start modifying it. 6 | 7 | import persistent 8 | import BTrees.OOBTree 9 | 10 | from stf.common.out import * 11 | from stf.common.abstracts import Module 12 | from stf.core.models import __groupofgroupofmodels__ 13 | from stf.core.dataset import __datasets__ 14 | from stf.core.notes import __notes__ 15 | from stf.core.connections import __group_of_group_of_connections__ 16 | from stf.core.models_constructors import __modelsconstructors__ 17 | from stf.core.labels import __group_of_labels__ 18 | from stf.core.database import __database__ 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | ################# 27 | ################# 28 | ################# 29 | class Template_Object(persistent.Persistent): 30 | """ This class is a template of a classic object. This is usually the place you want to do something. The new 'Object' that you want to store""" 31 | def __init__(self, id): 32 | self.id = id 33 | # This is an example dictionary of stuff that we want to store in the DB and make persistent. 34 | # self.dict_of_stuff = BTrees.OOBTree.BTree() 35 | 36 | def get_id(self): 37 | return self.id 38 | 39 | def set_id(self, id): 40 | self.id = id 41 | 42 | def set_name(self, name): 43 | self.name = name 44 | 45 | def get_name(self): 46 | return self.name 47 | 48 | # Mandatory __repr__ module. Something you want to identify each object with. Usefull for selecting objects later 49 | def __repr__(self): 50 | return('Id:' + str(self.get_id())) 51 | 52 | 53 | 54 | ###################### 55 | ###################### 56 | ###################### 57 | class Group_of_Template_Objects(Module, persistent.Persistent): 58 | """ The group of 'Objects' is only a structure to hold them together. Usefull to add them, delete them and general management """ 59 | ### Mandatory variables ### 60 | cmd = 'template_example_module' 61 | description = 'This module is a template to use for future modules. It stores permanently in the database the group of objects.' 62 | authors = ['Sebastian Garcia'] 63 | # Main dict of objects. The name of the attribute should be "main_dict" in this example 64 | main_dict = BTrees.OOBTree.BTree() 65 | ### End of Mandatory variables ### 66 | 67 | ### Mandatory Methods Don't change ### 68 | def __init__(self): 69 | # Call to our super init 70 | super(Group_of_Template_Objects, self).__init__() 71 | # Example of a parameter without arguments 72 | self.parser.add_argument('-l', '--list', action='store_true', help='List the objects in this group.') 73 | # Example of a parameter with arguments 74 | self.parser.add_argument('-p', '--printstate', metavar='printstate', help='Print some info about a specific object. Give the object id.') 75 | self.parser.add_argument('-g', '--generate', metavar='generate', help='Create a new object with a name. Give name.') 76 | 77 | def get_name(self): 78 | """ Return the name of the module""" 79 | return self.cmd 80 | 81 | # Mandatory Method! Don't change. 82 | def get_main_dict(self): 83 | """ Return the main dict where we store the info. Is going to the database""" 84 | return self.main_dict 85 | 86 | # Mandatory Method! Don't change. 87 | def set_main_dict(self, dict): 88 | """ Set the main dict where we store the info. From the database""" 89 | self.main_dict = dict 90 | ############ End of Mandatory Methods ######################### 91 | 92 | 93 | def get_object(self, id): 94 | return self.main_dict[id] 95 | 96 | def get_objects(self): 97 | return self.main_dict.values() 98 | 99 | def list_objects(self): 100 | print_info('List of objects') 101 | rows = [] 102 | for object in self.get_objects(): 103 | rows.append([ object.get_id(), object.get_name() ]) 104 | print(table(header=['Id', 'Name'], rows=rows)) 105 | 106 | def create_new_object(self, name): 107 | # Generate the new id 108 | try: 109 | new_id = self.main_dict[list(self.main_dict.keys())[-1]].get_id() + 1 110 | except (KeyError, IndexError): 111 | new_id = 1 112 | # Create the new object 113 | new_object = Template_Object(new_id) 114 | # Do something with it 115 | new_object.set_name(name) 116 | # Store on DB 117 | self.main_dict[new_id] = new_object 118 | 119 | 120 | 121 | # The run method runs every time that this command is used. Mandatory 122 | def run(self): 123 | ######### Mandatory part! don't delete ######################## 124 | # Register the structure in the database, so it is stored and use in the future. 125 | if not __database__.has_structure(Group_of_Template_Objects().get_name()): 126 | print_info('The structure is not registered.') 127 | __database__.set_new_structure(Group_of_Template_Objects()) 128 | else: 129 | main_dict = __database__.get_new_structure(Group_of_Template_Objects()) 130 | self.set_main_dict(main_dict) 131 | 132 | # List general help. Don't modify. 133 | def help(): 134 | self.log('info', self.description) 135 | 136 | # Run 137 | super(Group_of_Template_Objects, self).run() 138 | if self.args is None: 139 | return 140 | ######### End Mandatory part! ######################## 141 | 142 | 143 | # Process the command line and call the methods. Here add your own parameters 144 | if self.args.list: 145 | self.list_objects() 146 | elif self.args.generate: 147 | self.create_new_object(self.args.generate) 148 | else: 149 | print_error('At least one parameter is required in this module') 150 | self.usage() 151 | -------------------------------------------------------------------------------- /modules/visualize_1.py: -------------------------------------------------------------------------------- 1 | # Part of this file was taken from Viper - https://github.com/botherder/viper 2 | # The rest is from the Stratosphere Testing Framework 3 | # See the file 'LICENSE' for copying permission. 4 | 5 | # This is a template module showing how to create a module that has persistence in the database. To create your own command just copy this file and start modifying it. 6 | 7 | import persistent 8 | import BTrees.OOBTree 9 | import curses 10 | import multiprocessing 11 | from multiprocessing import Queue 12 | from multiprocessing import JoinableQueue 13 | import time 14 | from datetime import datetime 15 | import re 16 | 17 | 18 | from stf.common.out import * 19 | from stf.common.abstracts import Module 20 | from stf.core.dataset import __datasets__ 21 | from stf.core.models_constructors import __modelsconstructors__ 22 | from stf.core.database import __database__ 23 | from stf.core.connections import Flow 24 | from stf.core.models import Model 25 | 26 | 27 | ###################### 28 | ###################### 29 | ###################### 30 | class Screen(multiprocessing.Process): 31 | """ A class thread to run the screen """ 32 | def __init__(self, qscreen): 33 | multiprocessing.Process.__init__(self) 34 | self.qscreen = qscreen 35 | self.tuples = {} 36 | self.global_x_pos = 1 37 | self.y_min = 46 38 | self.f = open('log2','w') 39 | 40 | def get_tuple(self, tuple_id): 41 | """ Get a tuple, return its dict """ 42 | try: 43 | return self.tuples[tuple_id] 44 | except KeyError: 45 | # first time for this tuple 46 | self.tuples[tuple_id] = {} 47 | self.tuples[tuple_id]['y_pos'] = self.y_min 48 | self.tuples[tuple_id]['x_pos'] = self.global_x_pos 49 | (x_max, y_max) = self.screen.getmaxyx() 50 | if self.global_x_pos <= x_max - 2: 51 | self.global_x_pos += 1 52 | if 'tcp' in tuple_id.lower(): 53 | self.tuples[tuple_id]['color'] = curses.color_pair(1) 54 | elif 'udp' in tuple_id.lower(): 55 | self.tuples[tuple_id]['color'] = curses.color_pair(2) 56 | elif 'icmp' in tuple_id.lower(): 57 | self.tuples[tuple_id]['color'] = curses.color_pair(3) 58 | else: 59 | self.tuples[tuple_id]['color'] = curses.color_pair(4) 60 | # print the tuple 61 | self.screen.addstr(self.tuples[tuple_id]['x_pos'],0, tuple_id) 62 | return self.tuples[tuple_id] 63 | 64 | def run(self): 65 | try: 66 | while True: 67 | #print 'Is the queue empty?: {}'.format(self.qscreen.empty()) 68 | if not self.qscreen.empty(): 69 | order = self.qscreen.get() 70 | #self.logfile.write('Receiving the order'+order+'\n') 71 | if order == 'Start': 72 | stdscr = curses.initscr() 73 | curses.savetty() 74 | curses.start_color() 75 | curses.use_default_colors() 76 | self.screen = stdscr 77 | curses.init_pair(1, curses.COLOR_GREEN, -1) 78 | curses.init_pair(2, curses.COLOR_RED, -1) 79 | curses.init_pair(3, curses.COLOR_BLUE, -1) 80 | curses.init_pair(4, curses.COLOR_WHITE, -1) 81 | self.screen.bkgd(' ', curses.color_pair(1)) 82 | self.screen.bkgd(' ') 83 | curses.noecho() 84 | curses.cbreak() 85 | self.screen.keypad(1) 86 | # curses.curs_set. 0 means invisible cursor, 1 visible, 2 very visible 87 | curses.curs_set(0) 88 | self.screen.addstr(0,0, 'Live Stream', curses.A_BOLD) 89 | self.screen.refresh() 90 | self.qscreen.task_done() 91 | 92 | elif order == 'Stop': 93 | self.screen.addstr(0,0, 'Press any key to go back to stf screen... ', curses.A_BOLD) 94 | self.screen.getstr(0,0, 15) 95 | curses.nocbreak() 96 | self.screen.keypad(0) 97 | curses.echo() 98 | curses.curs_set(1) 99 | stdscr.refresh() 100 | # close 101 | self.qscreen.task_done() 102 | curses.resetty() 103 | self.f.close() 104 | return 105 | else: 106 | #self.screen.addstr(0,50, 'Receiving Data') 107 | self.screen.refresh() 108 | # Get the screen size 109 | (x_max, y_max) = self.screen.getmaxyx() 110 | # The order 111 | orig_tuple = order 112 | tuple_id = orig_tuple.get_id() 113 | # Get the amount of letters that fit on the screen 114 | state = orig_tuple.get_state()[-(y_max-self.y_min):] 115 | tuple = self.get_tuple(tuple_id) 116 | # Update the status bar 117 | if int(tuple['x_pos']) < x_max - 3: 118 | self.screen.addstr(0,20,tuple_id + " ", curses.A_BOLD) 119 | self.screen.refresh() 120 | #self.screen.addstr(int(tuple['x_pos']), int(tuple['y_pos']), state, tuple['color'] | curses.A_BOLD) 121 | #if int(tuple['x_pos']) <= xmax: 122 | self.screen.addstr(int(tuple['x_pos']), int(tuple['y_pos']), state, tuple['color']) 123 | #tuple['y_pos'] += len(state) 124 | self.screen.refresh() 125 | self.qscreen.task_done() 126 | except KeyboardInterrupt: 127 | curses.nocbreak() 128 | self.screen.keypad(0) 129 | curses.echo() 130 | curses.curs_set(1) 131 | curses.resetty() 132 | self.screen.refresh() 133 | except Exception as inst: 134 | print '\tProblem with Screen()' 135 | print type(inst) # the exception instance 136 | print inst.args # arguments stored in .args 137 | print inst # __str__ allows args to printed directly 138 | sys.exit(1) 139 | 140 | 141 | ###################### 142 | ###################### 143 | ###################### 144 | class Visualizations(Module): 145 | ### Mandatory variables ### 146 | cmd = 'visualize_1' 147 | description = 'This module visualize the connections of a dataset.' 148 | authors = ['Sebastian Garcia'] 149 | # Main dict of objects. The name of the attribute should be "main_dict" in this example 150 | main_dict = BTrees.OOBTree.BTree() 151 | ### End of Mandatory variables ### 152 | 153 | ### Mandatory Methods Don't change ### 154 | def __init__(self): 155 | # Call to our super init 156 | super(Visualizations, self).__init__() 157 | # Example of a parameter without arguments 158 | self.parser.add_argument('-v', '--visualize', metavar='datasetid', help='Visualize the connections in this dataset.') 159 | self.parser.add_argument('-x', '--multiplier', metavar='multiplier', type=float, default=0.0, help='Speed multiplier. 2 is twice, 3 is trice. -1 means to wait an equidistance but fake amount of time.') 160 | self.parser.add_argument('-f', '--filter', metavar='filter', nargs = '+', default="", help='Filter the connections to show. Keywords: name. Usage: name= name!=. Partial matching.') 161 | # A local list of models 162 | self.models = {} 163 | 164 | def get_name(self): 165 | """ Return the name of the module""" 166 | return self.cmd 167 | 168 | # Mandatory Method! Don't change. 169 | def get_main_dict(self): 170 | """ Return the main dict where we store the info. Is going to the database""" 171 | return self.main_dict 172 | 173 | # Mandatory Method! Don't change. 174 | def set_main_dict(self, dict): 175 | """ Set the main dict where we store the info. From the database""" 176 | self.main_dict = dict 177 | ############ End of Mandatory Methods ######################### 178 | 179 | 180 | def get_object(self, id): 181 | return self.main_dict[id] 182 | 183 | def get_objects(self): 184 | return self.main_dict.values() 185 | 186 | def visualize_dataset(self, dataset_id, multiplier, filter): 187 | # Get the netflow file 188 | self.dataset = __datasets__.get_dataset(dataset_id) 189 | try: 190 | self.binetflow_file = self.dataset.get_file_type('binetflow') 191 | except AttributeError: 192 | print_error('That testing dataset does no seem to exist.') 193 | return False 194 | # Open the file 195 | try: 196 | file = open(self.binetflow_file.get_name(), 'r') 197 | self.setup_screen() 198 | except IOError: 199 | print_error('The binetflow file is not present in the system.') 200 | return False 201 | # construct filter 202 | self.construct_filter(filter) 203 | # Clean the previous models from the constructor 204 | __modelsconstructors__.get_default_constructor().clean_models() 205 | # Remove the header 206 | header_line = file.readline().strip() 207 | # Find the separation character 208 | self.find_separator(header_line) 209 | # Extract the columns names 210 | self.find_columns_names(header_line) 211 | line = ','.join(file.readline().strip().split(',')[:14]) 212 | #logfile = open('log','w') 213 | while line: 214 | # Using our own extract_columns function makes this module more independent 215 | column_values = self.extract_columns_values(line) 216 | # Extract its 4-tuple. Find (or create) the tuple object 217 | tuple4 = column_values['SrcAddr']+'-'+column_values['DstAddr']+'-'+column_values['Dport']+'-'+column_values['Proto'] 218 | # Get the _local_ model. We don't want to mess with the real models in the database, but we need the structure to get the state 219 | model = self.get_model(tuple4) 220 | # filter 221 | if not self.apply_filter(tuple4): 222 | line = ','.join(file.readline().strip().split(',')[:14]) 223 | continue 224 | if not model: 225 | model = Model(tuple4) 226 | self.set_model(model) 227 | constructor_id = __modelsconstructors__.get_default_constructor().get_id() 228 | # Warning, here we depend on the modelsconstrutor 229 | model.set_constructor(__modelsconstructors__.get_constructor(constructor_id)) 230 | flow = Flow(0) # Fake flow id 231 | flow.add_starttime(column_values['StartTime']) 232 | flow.add_duration(column_values['Dur']) 233 | flow.add_proto(column_values['Proto']) 234 | flow.add_scraddr(column_values['SrcAddr']) 235 | flow.add_dir(column_values['Dir']) 236 | flow.add_dstaddr(column_values['DstAddr']) 237 | flow.add_dport(column_values['Dport']) 238 | flow.add_state(column_values['State']) 239 | flow.add_stos(column_values['sTos']) 240 | flow.add_dtos(column_values['dTos']) 241 | flow.add_totpkts(column_values['TotPkts']) 242 | flow.add_totbytes(column_values['TotBytes']) 243 | try: 244 | flow.add_srcbytes(column_values['SrcBytes']) 245 | except KeyError: 246 | # It can happen that we don't have the SrcBytes column 247 | pass 248 | try: 249 | flow.add_srcUdata(column_values['srcUdata']) 250 | except KeyError: 251 | # It can happen that we don't have the srcUdata column 252 | pass 253 | try: 254 | flow.add_dstUdata(column_values['dstUdata']) 255 | except KeyError: 256 | # It can happen that we don't have the dstUdata column 257 | pass 258 | try: 259 | flow.add_label(column_values['Label']) 260 | except KeyError: 261 | # It can happen that we don't have the label column 262 | pass 263 | # Add the flow 264 | model.add_flow(flow) 265 | # As fast as we can or with some delay? 266 | if multiplier != 0.0 and multiplier != -1: 267 | # Wait some time between flows. 268 | last_time = model.get_last_flow_time() 269 | current_time = datetime.strptime(column_values['StartTime'], '%Y/%m/%d %H:%M:%S.%f') 270 | if last_time: 271 | diff = current_time - last_time 272 | wait_time = diff.total_seconds() 273 | time.sleep(wait_time / multiplier) 274 | model.add_last_flow_time(current_time) 275 | # Wait the necessary time. After the visualization 276 | elif multiplier == -1: 277 | time.sleep(0.1) 278 | # Visualize this model 279 | self.qscreen.put(model) 280 | line = ','.join(file.readline().strip().split(',')[:14]) 281 | self.qscreen.put('Stop') 282 | file.close() 283 | #logfile.close() 284 | 285 | def set_model(self, model): 286 | self.models[model.get_id()] = model 287 | 288 | def clean_models(self): 289 | self.models = {} 290 | 291 | def get_model(self, tuple_id): 292 | try: 293 | return self.models[tuple_id] 294 | except KeyError: 295 | return False 296 | 297 | def find_separator(self, line): 298 | count_commas = len(line.split(',')) 299 | count_spaces = len(line.split(' ')) 300 | if count_commas >= count_spaces: 301 | self.line_separator = ',' 302 | elif count_spaces > count_commas: 303 | self.line_separator = ' ' 304 | else: 305 | self.line_separator = ',' 306 | 307 | def find_columns_names(self, line): 308 | """ Usually the columns in a binetflow file are 309 | StartTime,Dur,Proto,SrcAddr,Sport,Dir,DstAddr,Dport,State,sTos,dTos,TotPkts,TotBytes,SrcBytes,srcUdata,dstUdata,Label 310 | """ 311 | self.columns_names = line.split(self.line_separator) 312 | 313 | def extract_columns_values(self, line): 314 | """ Given a line text of a flow, extract the values for each column. The main difference with this function and the one in connections.py is that we don't use the src and dst data. """ 315 | column_values = {} 316 | i = 0 317 | original_values = line.split(self.line_separator) 318 | temp_values = original_values 319 | if len(temp_values) > len(self.columns_names): 320 | # If there is only one occurrence of the separator char, then try to recover... 321 | # Find the index of srcudata 322 | srcudata_index_starts = 0 323 | for values in temp_values: 324 | if 's[' in values: 325 | break 326 | else: 327 | srcudata_index_starts += 1 328 | # Find the index of dstudata 329 | dstudata_index_starts = 0 330 | for values in temp_values: 331 | if 'd[' in values: 332 | break 333 | else: 334 | dstudata_index_starts += 1 335 | # Get all the src data 336 | srcudata_index_ends = dstudata_index_starts 337 | temp_srcudata = temp_values[srcudata_index_starts:srcudata_index_ends] 338 | srcudata = '' 339 | for i in temp_srcudata: 340 | srcudata = srcudata + i 341 | # Get all the dst data. The end is one before the last field. That we know is the label. 342 | dstudata_index_ends = len(temp_values) - 1 343 | temp_dstudata = temp_values[dstudata_index_starts:dstudata_index_ends] 344 | dstudata = '' 345 | for j in temp_dstudata: 346 | dstudata = dstudata + j 347 | label = temp_values[-1] 348 | end_of_good_data = srcudata_index_starts 349 | # Rewrite temp_values 350 | temp_values = temp_values[0:end_of_good_data] 351 | temp_values.append(srcudata) 352 | temp_values.append(dstudata) 353 | temp_values.append(label) 354 | index = 0 355 | try: 356 | for value in temp_values: 357 | column_values[self.columns_names[index]] = value 358 | index += 1 359 | except IndexError: 360 | # Even with our fix, some data still has problems. Usually it means that there is no src data being sent, so we can not find the start of the data. 361 | print_error('There was some error reading the data inside a flow. Most surely it includes the field separator of the flows. We will keep the flow, but not its data.') 362 | # Just get the normal flow fields 363 | index = 0 364 | for value in temp_values: 365 | if index <= 13: 366 | column_values[self.columns_names[index]] = value 367 | index += 1 368 | else: 369 | break 370 | column_values['srcUdata'] = 'Deleted because of inconsistencies' 371 | column_values['dstUdata'] = 'Deleted because of inconsistencies' 372 | column_values['Label'] = original_values[-1] 373 | return column_values 374 | 375 | def construct_filter(self, filter): 376 | """ Get the filter string and decode all the operations """ 377 | # If the filter string is empty, delete the filter variable 378 | if not filter: 379 | try: 380 | del self.filter 381 | except: 382 | pass 383 | return True 384 | self.filter = [] 385 | # Get the individual parts. We only support and's now. 386 | for part in filter: 387 | # Get the key 388 | try: 389 | key = re.split('<|>|=|\!=', part)[0] 390 | value = re.split('<|>|=|\!=', part)[1] 391 | except IndexError: 392 | # No < or > or = or != in the string. Just stop. 393 | break 394 | try: 395 | part.index('<') 396 | operator = '<' 397 | except ValueError: 398 | pass 399 | try: 400 | part.index('>') 401 | operator = '>' 402 | except ValueError: 403 | pass 404 | # We should search for != before = 405 | try: 406 | part.index('!=') 407 | operator = '!=' 408 | except ValueError: 409 | # Now we search for = 410 | try: 411 | part.index('=') 412 | operator = '=' 413 | except ValueError: 414 | pass 415 | self.filter.append((key, operator, value)) 416 | 417 | def apply_filter(self, tuple4): 418 | """ Use the stored filter to know what we should match""" 419 | responses = [] 420 | try: 421 | self.filter 422 | except AttributeError: 423 | # If we don't have any filter string, just return true and show everything 424 | return True 425 | # Check each filter 426 | for filter in self.filter: 427 | key = filter[0] 428 | operator = filter[1] 429 | value = filter[2] 430 | if key == 'name': 431 | # For filtering based on the label assigned to the model with stf (contrary to the flow label) 432 | if operator == '=': 433 | if value in tuple4: 434 | responses.append(True) 435 | else: 436 | responses.append(False) 437 | elif operator == '!=': 438 | if value not in tuple4: 439 | responses.append(True) 440 | else: 441 | responses.append(False) 442 | else: 443 | return False 444 | for response in responses: 445 | if not response: 446 | return False 447 | return True 448 | 449 | def setup_screen(self): 450 | # Create the queue 451 | self.qscreen = JoinableQueue() 452 | # Create the thread 453 | self.screenThread = Screen(self.qscreen) 454 | self.screenThread.start() 455 | # Waint until the screen is initialized 456 | # First send the message 457 | self.qscreen.put('Start') 458 | # Then wait for an answer 459 | self.qscreen.join() 460 | 461 | # The run method runs every time that this command is used. Mandatory 462 | def run(self): 463 | ######### Mandatory part! don't delete ######################## 464 | # Register the structure in the database, so it is stored and use in the future. 465 | if not __database__.has_structure(Visualizations().get_name()): 466 | print_info('The structure is not registered.') 467 | __database__.set_new_structure(Visualizations()) 468 | else: 469 | main_dict = __database__.get_new_structure(Visualizations()) 470 | self.set_main_dict(main_dict) 471 | 472 | # List general help. Don't modify. 473 | def help(): 474 | self.log('info', self.description) 475 | 476 | # Run 477 | super(Visualizations, self).run() 478 | if self.args is None: 479 | return 480 | ######### End Mandatory part! ######################## 481 | 482 | 483 | # Process the command line and call the methods. Here add your own parameters 484 | if self.args.visualize: 485 | self.visualize_dataset(int(self.args.visualize), self.args.multiplier, self.args.filter) 486 | else: 487 | print_error('At least one of the parameter is required in this module') 488 | self.usage() 489 | -------------------------------------------------------------------------------- /stf.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # This file is part of the Stratosphere Testing Framework 3 | # See the file 'LICENSE' for copying permission. 4 | 5 | import argparse 6 | import os 7 | 8 | from stf.core.ui import console 9 | from stf.core.configuration import __configuration__ 10 | 11 | parser = argparse.ArgumentParser() 12 | parser.add_argument('-c', '--config', help='Configuration file.', action='store', required=False, type=str) 13 | args = parser.parse_args() 14 | 15 | # Read the conf file. 16 | # If we were given a conf file, use it 17 | if args.config and os.path.isfile(args.config): 18 | config_file = args.config 19 | # If not, search for the conf file in our local folder 20 | elif os.path.isfile('./confs/stf.conf'): 21 | config_file = './confs/stf.conf' 22 | else: 23 | print 'No configuration file found. Either give one with -c or put one in the local confs folder.' 24 | exit(-1) 25 | 26 | # Load the configuration 27 | if __configuration__.read_conf_file(os.path.expanduser(config_file)): 28 | # Create the console and start 29 | c = console.Console() 30 | c.start() 31 | 32 | -------------------------------------------------------------------------------- /stf/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/stratosphereips/StratosphereTestingFramework/381eca74358dfa3847727314f2ae615cd5c8170a/stf/__init__.py -------------------------------------------------------------------------------- /stf/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/stratosphereips/StratosphereTestingFramework/381eca74358dfa3847727314f2ae615cd5c8170a/stf/__init__.pyc -------------------------------------------------------------------------------- /stf/common/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/stratosphereips/StratosphereTestingFramework/381eca74358dfa3847727314f2ae615cd5c8170a/stf/common/__init__.py -------------------------------------------------------------------------------- /stf/common/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/stratosphereips/StratosphereTestingFramework/381eca74358dfa3847727314f2ae615cd5c8170a/stf/common/__init__.pyc -------------------------------------------------------------------------------- /stf/common/abstracts.py: -------------------------------------------------------------------------------- 1 | # This file is part of Viper - https://github.com/botherder/viper 2 | # See the file 'LICENSE' for copying permission. 3 | 4 | import argparse 5 | 6 | class ArgumentErrorCallback(Exception): 7 | def __init__(self, message, level=''): 8 | self.message = message.strip() + '\n' 9 | self.level = level.strip() 10 | def __str__(self): 11 | return '{}: {}'.format(self.level, self.message) 12 | def get(self): 13 | return self.level, self.message 14 | 15 | 16 | class ArgumentParser(argparse.ArgumentParser): 17 | def print_usage(self): 18 | raise ArgumentErrorCallback(self.format_usage()) 19 | def print_help(self): 20 | raise ArgumentErrorCallback(self.format_help()) 21 | def error(self, message): 22 | raise ArgumentErrorCallback(message, 'error') 23 | def exit(self, status, message=None): 24 | if message is not None: 25 | raise ArgumentErrorCallback(message) 26 | 27 | 28 | class Module(object): 29 | cmd = '' 30 | description = '' 31 | command_line = [] 32 | args = None 33 | authors = [] 34 | output = [] 35 | 36 | def __init__(self): 37 | self.parser = ArgumentParser(prog=self.cmd, description=self.description) 38 | 39 | def set_commandline(self, command): 40 | self.command_line = command 41 | 42 | def log(self, event_type, event_data): 43 | self.output.append(dict( 44 | type=event_type, 45 | data=event_data 46 | )) 47 | 48 | def usage(self): 49 | self.log('', self.parser.format_usage()) 50 | 51 | def help(self): 52 | self.log('', self.parser.format_help()) 53 | 54 | def run(self): 55 | try: 56 | self.args = self.parser.parse_args(self.command_line) 57 | except ArgumentErrorCallback as e: 58 | self.log(*e.get()) 59 | -------------------------------------------------------------------------------- /stf/common/ap.py: -------------------------------------------------------------------------------- 1 | """ 2 | Package that allows you to plot simple graphs in ASCII, a la matplotlib. 3 | 4 | This package is a inspired from Imri Goldberg's ASCII-Plotter 1.0 5 | (https://pypi.python.org/pypi/ASCII-Plotter/1.0) 6 | 7 | At a time I was enoyed by security not giving me direct access to my computer, 8 | and thus to quickly make figures from python, I looked at how I could make 9 | quick and dirty ASCII figures. But if I were to develop something, I wanted 10 | something that can be used with just python and possible standard-ish packages 11 | (numpy, scipy). 12 | 13 | So I came up with this package after many iterations based of ASCII-plotter. 14 | I added the feature to show multiple curves on one plot with different markers. 15 | And I also made the usage, close to matplotlib, such that there is a plot, 16 | hist, hist2d and imshow functions. 17 | 18 | 19 | TODO: 20 | imshow does not plot axis yet. 21 | make a correct documentation 22 | """ 23 | import math as _math 24 | import numpy as np 25 | 26 | 27 | __version__ = 0.9 28 | __author__ = 'M. Fouesneau' 29 | 30 | __all__ = ['markers', 'ACanvas', 'AData', 'AFigure', 31 | 'hist', 'hist2d', 'imshow', 'percentile_imshow', 32 | 'plot', 'stem', 'stemify', 'step', 'steppify', 33 | '__version__', '__author__'] 34 | 35 | 36 | markers = { '-' : u'None' , # solid line style 37 | ',': u'\u2219', # point marker 38 | '.': u'\u2218', # pixel marker 39 | '.f': u'\u2218', # pixel marker 40 | 'o': u'\u25CB', # circle marker 41 | 'of': u'\u25CF', # circle marker 42 | 'v': u'\u25BD', # triangle_down marker 43 | 'vf': u'\u25BC', # filler triangle_down marker 44 | '^': u'\u25B3', # triangle_up marker 45 | '^f': u'\u25B2', # filled triangle_up marker 46 | '<': u'\u25C1', # triangle_left marker 47 | '': u'\u25B7', # triangle_right marker 49 | '>f': u'\u25B6', # filled triangle_right marker 50 | 's': u'\u25FD', # square marker 51 | 'sf': u'\u25FC', # square marker 52 | '*': u'\u2606', # star marker 53 | '*f': u'\u2605', # star marker 54 | '+': u'\u271A', # plus marker 55 | 'x': u'\u274C', # x marker 56 | 'd': u'\u25C7', # diamond marker 57 | 'df': u'\u25C6' # filled diamond marker 58 | } 59 | 60 | 61 | def _sign(x): 62 | """ Return the sign of x 63 | INPUTS 64 | ------ 65 | x: number 66 | value to get the sign of 67 | 68 | OUTPUTS 69 | ------- 70 | s: signed int 71 | -1, 0 or 1 if negative, null or positive 72 | """ 73 | 74 | if (x > 0): 75 | return 1 76 | elif (x == 0): 77 | return 0 78 | else: 79 | return -1 80 | 81 | 82 | def _transpose(mat): 83 | """ Transpose matrice made of lists 84 | INPUTS 85 | ------ 86 | mat: iterable 2d list like 87 | 88 | OUTPUTS 89 | ------- 90 | r: list of list, 2d list like 91 | transposed matrice 92 | """ 93 | r = [ [x[i] for x in mat] for i in range(len(mat[0])) ] 94 | return r 95 | 96 | 97 | def _y_reverse(mat): 98 | """ Reverse the y axis of a 2d list-like 99 | INPUTS 100 | ------ 101 | mat: list of lists 102 | the matrix to reverse on axis 0 103 | OUTPUTS 104 | ------- 105 | r: list of lists 106 | the reversed version 107 | """ 108 | r = [ list(reversed(mat_i)) for mat_i in mat ] 109 | return r 110 | 111 | 112 | class AData(object): 113 | """ Data container for ascii AFigure """ 114 | def __init__(self, x, y, marker='_.', plot_slope=True): 115 | """ Constructor 116 | INPUTS 117 | ------ 118 | x: iterable 119 | x values 120 | y: iterable 121 | y values 122 | 123 | KEYWORDS 124 | -------- 125 | marker: str 126 | marker for the data. 127 | if None or empty, the curve will be plotted 128 | if the first character of the marker is '_' then unicode markers will be called: 129 | 130 | marker repr description 131 | ======== =========== ============================= 132 | '-' u'None' solid line style 133 | ',' u'\u2219' point marker 134 | '.' u'\u2218' pixel marker 135 | '.f' u'\u2218' pixel marker 136 | 'o' u'\u25CB' circle marker 137 | 'of' u'\u25CF' circle marker 138 | 'v' u'\u25BD' triangle_down marker 139 | 'vf' u'\u25BC' filler triangle_down marker 140 | '^' u'\u25B3' triangle_up marker 141 | '^f' u'\u25B2' filled triangle_up marker 142 | '<' u'\u25C1' triangle_left marker 143 | '' u'\u25B7' triangle_right marker 145 | '>f' u'\u25B6' filled triangle_right marker 146 | 's' u'\u25FD' square marker 147 | 'sf' u'\u25FC' square marker 148 | '*' u'\u2606' star marker 149 | '*f' u'\u2605' star marker 150 | '+' u'\u271A' plus marker 151 | 'x' u'\u274C' x marker 152 | 'd' u'\u25C7' diamond marker 153 | 'df' u'\u25C6' filled diamond marker 154 | 155 | plot_slope: bool 156 | if set, the curve will be plotted 157 | """ 158 | self.x = x 159 | self.y = y 160 | self.plot_slope = plot_slope 161 | self.set_marker(marker) 162 | 163 | def set_marker(self, marker): 164 | """ set the marker of the data 165 | INPUTS 166 | ------ 167 | marker: str 168 | marker for the data. 169 | see constructor for marker descriptions 170 | """ 171 | if marker in [None, 'None', u'None', '']: 172 | self.plot_slope = True 173 | self.marker = '' 174 | elif marker[0] == '_': 175 | self.marker = markers[marker[1:]] 176 | else: 177 | self.marker = marker 178 | 179 | def extent(self): 180 | """ return the extention of the data 181 | OUPUTS 182 | ------ 183 | e: list 184 | [ min(x), max(x), min(y), max(y) ] 185 | """ 186 | return [min(self.x), max(self.x), min(self.y), max(self.y)] 187 | 188 | def __repr__(self): 189 | s = 'AData: %s\n' % object.__repr__(self) 190 | return s 191 | 192 | 193 | class ACanvas(object): 194 | """ Canvas of a AFigure instance. A Canvas handles all transformations 195 | between data space and figure space accounting for scaling and pixels 196 | 197 | In general there is no need to access the canvas directly 198 | """ 199 | def __init__(self, shape=None, margins=None, xlim=None, ylim=None): 200 | """ Constructor 201 | KEYWORDS 202 | -------- 203 | shape: tuple of 2 ints 204 | shape of the canvas in number of characters: (width, height) 205 | 206 | margins: tuple of 2 floats 207 | fractional margins 208 | 209 | xlim: tuple of 2 floats 210 | limits of the xaxis 211 | 212 | ylim: tuple of 2 floats 213 | limits of the yaxis 214 | """ 215 | self.shape = shape or (50, 20) 216 | self.margins = margins or (0.05, 0.1) 217 | self._xlim = xlim or [0, 1] 218 | self._ylim = ylim or [0, 1] 219 | self.auto_adjust = True 220 | self.margin_factor = 1 221 | 222 | @property 223 | def x_size(self): 224 | """ return the width """ 225 | return self.shape[0] 226 | 227 | @property 228 | def y_size(self): 229 | """ return the height """ 230 | return self.shape[1] 231 | 232 | @property 233 | def x_margin(self): 234 | """ return the margin in x """ 235 | return self.margins[0] 236 | 237 | @property 238 | def y_margin(self): 239 | """ return the margin in y """ 240 | return self.margins[1] 241 | 242 | def xlim(self, vmin=None, vmax=None): 243 | """ 244 | Get or set the *x* limits of the current axes. 245 | 246 | KEYWORDS 247 | -------- 248 | vmin: float 249 | lower limit 250 | vmax: float 251 | upper limit 252 | 253 | xmin, xmax = xlim() # return the current xlim 254 | xlim( (xmin, xmax) ) # set the xlim to xmin, xmax 255 | xlim( xmin, xmax ) # set the xlim to xmin, xmax 256 | """ 257 | if vmin is None and vmax is None: 258 | return self._xlim 259 | elif hasattr(vmin, '__iter__'): 260 | self._xlim = vmin[:2] 261 | else: 262 | self._xlim = [vmin, vmax] 263 | if self._xlim[0] == self._xlim[1]: 264 | self._xlim[1] += 1 265 | 266 | self._xlim[0] -= self.x_mod 267 | self._xlim[1] += self.x_mod 268 | 269 | def ylim(self, vmin=None, vmax=None): 270 | """ 271 | Get or set the *y* limits of the current axes. 272 | 273 | KEYWORDS 274 | -------- 275 | vmin: float 276 | lower limit 277 | vmax: float 278 | upper limit 279 | 280 | ymin, ymax = ylim() # return the current xlim 281 | ylim( (ymin, ymax) ) # set the xlim to xmin, xmax 282 | ylim( ymin, ymax ) # set the xlim to xmin, xmax 283 | """ 284 | if vmin is None and vmax is None: 285 | return self._ylim 286 | elif hasattr(vmin, '__iter__'): 287 | self._ylim = vmin[:2] 288 | else: 289 | self._ylim = [vmin, vmax] 290 | if self._ylim[0] == self._ylim[1]: 291 | self._ylim[1] += 1 292 | 293 | self._ylim[0] -= self.y_mod 294 | self._ylim[1] += self.y_mod 295 | 296 | @property 297 | def min_x(self): 298 | """ return the lower x limit """ 299 | return self._xlim[0] 300 | 301 | @property 302 | def max_x(self): 303 | """ return the upper x limit """ 304 | return self._xlim[1] 305 | 306 | @property 307 | def min_y(self): 308 | """ return the lower y limit """ 309 | return self._ylim[0] 310 | 311 | @property 312 | def max_y(self): 313 | """ return the upper y limit """ 314 | return self._ylim[1] 315 | 316 | @property 317 | def x_step(self): 318 | return float(self.max_x - self.min_x) / float(self.x_size) 319 | 320 | @property 321 | def y_step(self): 322 | return float(self.max_y - self.min_y) / float(self.y_size) 323 | 324 | @property 325 | def ratio(self): 326 | return self.y_step / self.x_step 327 | 328 | @property 329 | def x_mod(self): 330 | return (self.max_x - self.min_x) * self.x_margin 331 | 332 | @property 333 | def y_mod(self): 334 | return (self.max_y - self.min_y) * self.y_margin 335 | 336 | def extent(self, margin_factor=None): 337 | margin_factor = margin_factor or self.margin_factor 338 | min_x = (self.min_x + self.x_mod * margin_factor) 339 | max_x = (self.max_x - self.x_mod * margin_factor) 340 | min_y = (self.min_y + self.y_mod * margin_factor) 341 | max_y = (self.max_y - self.y_mod * margin_factor) 342 | return (min_x, max_x, min_y, max_y) 343 | 344 | def extent_str(self, margin=None): 345 | 346 | def transform(val, fmt): 347 | if abs(val) < 1: 348 | _str = "%+.2g" % val 349 | elif fmt is not None: 350 | _str = fmt % val 351 | else: 352 | _str = None 353 | return _str 354 | 355 | e = self.extent(margin) 356 | 357 | xfmt = self.x_str() 358 | yfmt = self.y_str() 359 | return transform(e[0], xfmt), transform(e[1], xfmt), transform(e[2], yfmt), transform(e[3], yfmt) 360 | 361 | def x_str(self): 362 | if self.x_size < 16: 363 | x_str = None 364 | elif self.x_size < 23: 365 | x_str = "%+.2g" 366 | else: 367 | x_str = "%+g" 368 | return x_str 369 | 370 | def y_str(self): 371 | if self.x_size < 8: 372 | y_str = None 373 | elif self.x_size < 11: 374 | y_str = "%+.2g" 375 | else: 376 | y_str = "%+g" 377 | return y_str 378 | 379 | def coords_inside_buffer(self, x, y): 380 | return (0 <= x < self.x_size) and (0 < y < self.y_size) 381 | 382 | def coords_inside_data(self, x, y): 383 | """ return if (x,y) covered by the data box 384 | x, y: float 385 | coordinates to test 386 | """ 387 | return (self.min_x <= x < self.max_x) and (self.min_y <= y < self.max_y) 388 | 389 | def _clip_line(self, line_pt_1, line_pt_2): 390 | """ clip a line to the canvas """ 391 | 392 | e = self.extent() 393 | x_min = min(line_pt_1[0], line_pt_2[0]) 394 | x_max = max(line_pt_1[0], line_pt_2[0]) 395 | y_min = min(line_pt_1[1], line_pt_2[1]) 396 | y_max = max(line_pt_1[1], line_pt_2[1]) 397 | 398 | if line_pt_1[0] == line_pt_2[0]: 399 | return ( ( line_pt_1[0], max(y_min, e[1]) ), 400 | ( line_pt_1[0], min(y_max, e[3]) )) 401 | 402 | if line_pt_1[1] == line_pt_2[1]: 403 | return ( ( max(x_min, e[0]), line_pt_1[1] ), 404 | ( min(x_max, e[2]), line_pt_1[1] )) 405 | 406 | if ( (e[0] <= line_pt_1[0] < e[2]) and 407 | (e[1] <= line_pt_1[1] < e[3]) and 408 | (e[0] <= line_pt_2[0] < e[2]) and 409 | (e[1] <= line_pt_2[1] < e[3]) ): 410 | return line_pt_1, line_pt_2 411 | 412 | ts = [0.0, 413 | 1.0, 414 | float(e[0] - line_pt_1[0]) / (line_pt_2[0] - line_pt_1[0]), 415 | float(e[2] - line_pt_1[0]) / (line_pt_2[0] - line_pt_1[0]), 416 | float(e[1] - line_pt_1[1]) / (line_pt_2[1] - line_pt_1[1]), 417 | float(e[3] - line_pt_1[1]) / (line_pt_2[1] - line_pt_1[1]) 418 | ] 419 | ts.sort() 420 | 421 | if (ts[2] < 0) or (ts[2] >= 1) or (ts[3] < 0) or (ts[2] >= 1): 422 | return None 423 | 424 | result = [(pt_1 + t * (pt_2 - pt_1)) for t in (ts[2], ts[3]) for (pt_1, pt_2) in zip(line_pt_1, line_pt_2)] 425 | 426 | return ( result[:2], result[2:] ) 427 | 428 | 429 | class AFigure(object): 430 | 431 | def __init__(self, shape=(80, 20), margins=(0.05, 0.1), draw_axes=True, newline='\n', 432 | plot_labels=True, xlim=None, ylim=None, **kwargs): 433 | 434 | self.canvas = ACanvas(shape, margins=margins, xlim=xlim, ylim=ylim) 435 | self.draw_axes = draw_axes 436 | self.new_line = newline 437 | self.plot_labels = plot_labels 438 | self.output_buffer = None 439 | self.tickSymbols = u'\u253C' # "+" 440 | self.x_axis_symbol = u'\u2500' # u"\u23bc" # "-" 441 | self.y_axis_symbol = u'\u2502' # "|" 442 | self.data = [] 443 | 444 | def xlim(self, vmin=None, vmax=None): 445 | return self.canvas.xlim(vmin, vmax) 446 | 447 | def ylim(self, vmin=None, vmax=None): 448 | return self.canvas.ylim(vmin, vmax) 449 | 450 | def get_coord(self, val, min, step, limits=None): 451 | result = int((val - min) / step) 452 | if limits is not None: 453 | if result <= limits[0]: 454 | result = limits[0] 455 | elif result >= limits[1]: 456 | result = limits[1] - 1 457 | return result 458 | 459 | def _draw_axes(self): 460 | zero_x = self.get_coord(0, self.canvas.min_x, self.canvas.x_step, limits=[1, self.canvas.x_size]) 461 | if zero_x >= self.canvas.x_size: 462 | zero_x = self.canvas.x_size - 1 463 | for y in xrange(self.canvas.y_size): 464 | self.output_buffer[zero_x][y] = self.y_axis_symbol 465 | 466 | zero_y = self.get_coord(0, self.canvas.min_y, self.canvas.y_step, limits=[1, self.canvas.y_size]) 467 | if zero_y >= self.canvas.y_size: 468 | zero_y = self.canvas.y_size - 1 469 | for x in xrange(self.canvas.x_size): 470 | self.output_buffer[x][zero_y] = self.x_axis_symbol # u'\u23bc' 471 | 472 | self.output_buffer[zero_x][zero_y] = self.tickSymbols # "+" 473 | 474 | def _get_symbol_by_slope(self, slope, default_symbol): 475 | """ Return a line oriented directed approximatively along the slope value """ 476 | if slope > _math.tan(3 * _math.pi / 8): 477 | draw_symbol = "|" 478 | elif _math.tan(_math.pi / 8) < slope < _math.tan(3 * _math.pi / 8): 479 | draw_symbol = u'\u27cb' # "/" 480 | elif abs(slope) < _math.tan(_math.pi / 8): 481 | draw_symbol = "-" 482 | elif slope < _math.tan(-_math.pi / 8) and slope > _math.tan(-3 * _math.pi / 8): 483 | draw_symbol = u'\u27CD' # "\\" 484 | elif slope < _math.tan(-3 * _math.pi / 8): 485 | draw_symbol = "|" 486 | else: 487 | draw_symbol = default_symbol 488 | return draw_symbol 489 | 490 | def _plot_labels(self): 491 | 492 | if self.canvas.y_size < 2: 493 | return 494 | 495 | act_min_x, act_max_x, act_min_y, act_max_y = self.canvas.extent() 496 | 497 | min_x_coord = self.get_coord(act_min_x, self.canvas.min_x, self.canvas.x_step, limits=[0, self.canvas.x_size]) 498 | max_x_coord = self.get_coord(act_max_x, self.canvas.min_x, self.canvas.x_step, limits=[0, self.canvas.x_size]) 499 | min_y_coord = self.get_coord(act_min_y, self.canvas.min_y, self.canvas.y_step, limits=[1, self.canvas.y_size]) 500 | max_y_coord = self.get_coord(act_max_y, self.canvas.min_y, self.canvas.y_step, limits=[1, self.canvas.y_size]) 501 | 502 | x_zero_coord = self.get_coord(0, self.canvas.min_x, self.canvas.x_step, limits=[0, self.canvas.x_size]) 503 | y_zero_coord = self.get_coord(0, self.canvas.min_y, self.canvas.y_step, limits=[1, self.canvas.y_size]) 504 | 505 | self.output_buffer[x_zero_coord][min_y_coord] = self.tickSymbols 506 | self.output_buffer[x_zero_coord][max_y_coord] = self.tickSymbols 507 | self.output_buffer[min_x_coord][y_zero_coord] = self.tickSymbols 508 | self.output_buffer[max_x_coord][y_zero_coord] = self.tickSymbols 509 | 510 | min_x_str, max_x_str, min_y_str, max_y_str = self.canvas.extent_str() 511 | if (self.canvas.x_str() is not None): 512 | for i, c in enumerate(min_x_str): 513 | self.output_buffer[min_x_coord + i + 1][y_zero_coord - 1] = c 514 | for i, c in enumerate(max_x_str): 515 | self.output_buffer[max_x_coord + i - len(max_x_str)][y_zero_coord - 1] = c 516 | 517 | if (self.canvas.y_str() is not None): 518 | for i, c in enumerate(max_y_str): 519 | self.output_buffer[x_zero_coord + i + 1][max_y_coord] = c 520 | for i, c in enumerate(min_y_str): 521 | self.output_buffer[x_zero_coord + i + 1][min_y_coord] = c 522 | 523 | def _plot_line(self, start, end, data): 524 | """ plot a line from start = (x0, y0) to end = (x1, y1) """ 525 | 526 | clipped_line = self.canvas._clip_line(start, end) 527 | 528 | if clipped_line is None: 529 | return False 530 | 531 | start, end = clipped_line 532 | 533 | x0 = self.get_coord(start[0], self.canvas.min_x, self.canvas.x_step) 534 | y0 = self.get_coord(start[1], self.canvas.min_y, self.canvas.y_step) 535 | x1 = self.get_coord(end[0], self.canvas.min_x, self.canvas.x_step) 536 | y1 = self.get_coord(end[1], self.canvas.min_y, self.canvas.y_step) 537 | 538 | if (x0, y0) == (x1, y1): 539 | return True 540 | 541 | #x_zero_coord = self.get_coord(0, self.canvas.min_x, self.canvas.x_step) 542 | y_zero_coord = self.get_coord(0, self.canvas.min_y, self.canvas.y_step, limits=[1, self.canvas.y_size]) 543 | 544 | if start[0] - end[0] == 0: 545 | draw_symbol = u"|" 546 | elif start[1] - end[1] == 0: 547 | draw_symbol = '-' 548 | else: 549 | slope = (1.0 / self.canvas.ratio) * (end[1] - start[1]) / (end[0] - start[0]) 550 | draw_symbol = self._get_symbol_by_slope(slope, data.marker) 551 | 552 | dx = x1 - x0 553 | dy = y1 - y0 554 | if abs(dx) > abs(dy): 555 | s = _sign(dx) 556 | slope = float(dy) / dx 557 | for i in range(0, abs(int(dx))): 558 | cur_draw_symbol = draw_symbol 559 | x = i * s 560 | cur_y = int(y0 + slope * x) 561 | if (self.draw_axes) and (cur_y == y_zero_coord) and (draw_symbol == self.x_axis_symbol): 562 | cur_draw_symbol = "-" 563 | self.output_buffer[x0 + x][cur_y] = cur_draw_symbol 564 | else: 565 | s = _sign(dy) 566 | slope = float(dx) / dy 567 | for i in range(0, abs(int(dy))): 568 | y = i * s 569 | cur_draw_symbol = draw_symbol 570 | cur_y = y0 + y 571 | if (self.draw_axes) and (cur_y == y_zero_coord) and (draw_symbol == self.x_axis_symbol): 572 | cur_draw_symbol = "-" 573 | self.output_buffer[int(x0 + slope * y)][cur_y] = cur_draw_symbol 574 | 575 | return False 576 | 577 | def _plot_data_with_slope(self, data): 578 | xy = list(zip(data.x, data.y)) 579 | 580 | #sort according to the x coord 581 | xy.sort(key=lambda c: c[0]) 582 | prev_p = xy[0] 583 | e_xy = enumerate(xy) 584 | e_xy.next() 585 | for i, (xi, yi) in e_xy: 586 | line = self._plot_line(prev_p, (xi, yi), data) 587 | prev_p = (xi, yi) 588 | 589 | # if no line, then symbol 590 | if not line & self.canvas.coords_inside_data(xi, yi): 591 | draw_symbol = data.marker 592 | 593 | px, py = xy[i - 1] 594 | nx, ny = xy[i] 595 | 596 | if abs(nx - px) > 0.000001: 597 | slope = (1.0 / self.canvas.ratio) * (ny - py) / (nx - px) 598 | draw_symbol = self._get_symbol_by_slope(slope, draw_symbol) 599 | 600 | x_coord = self.get_coord(xi, self.canvas.min_x, self.canvas.x_step) 601 | y_coord = self.get_coord(yi, self.canvas.min_y, self.canvas.y_step) 602 | 603 | if self.canvas.coords_inside_buffer(x_coord, y_coord): 604 | y0_coord = self.get_coord(0, self.canvas.min_y, self.canvas.y_step) 605 | if self.draw_axes: 606 | if (y_coord == y0_coord) and (draw_symbol == u"\u23bc"): 607 | draw_symbol = "=" 608 | self.output_buffer[x_coord][y_coord] = draw_symbol 609 | 610 | def _plot_data(self, data): 611 | if data.plot_slope: 612 | self._plot_data_with_slope(data) 613 | else: 614 | for x, y in zip(data.x, data.y): 615 | if self.canvas.coords_inside_data(x, y): 616 | x_coord = self.get_coord(x, self.canvas.min_x, self.canvas.x_step) 617 | y_coord = self.get_coord(y, self.canvas.min_y, self.canvas.y_step) 618 | 619 | if self.canvas.coords_inside_buffer(x_coord, y_coord): 620 | self.output_buffer[x_coord][y_coord] = data.marker 621 | 622 | def auto_limits(self): 623 | if self.canvas.auto_adjust is True: 624 | min_x = 0. 625 | max_x = 0. 626 | min_y = 0. 627 | max_y = 0. 628 | for dk in self.data: 629 | ek = dk.extent() 630 | min_x = min(min_x, min(ek[:2])) 631 | min_y = min(min_y, min(ek[2:])) 632 | max_x = max(max_x, max(ek[:2])) 633 | max_y = max(max_y, max(ek[2:])) 634 | self.canvas.xlim(min_x, max_x) 635 | self.canvas.ylim(min_y, max_y) 636 | 637 | def append_data(self, data): 638 | self.data.append(data) 639 | self.auto_limits() 640 | 641 | def plot(self, x_seq, y_seq=None, marker=None, plot_slope=False, xlim=None, ylim=None): 642 | 643 | if y_seq is None: 644 | y_seq = x_seq[:] 645 | x_seq = range(len(y_seq)) 646 | 647 | data = AData(x_seq, y_seq, marker=marker, plot_slope=plot_slope) 648 | self.append_data(data) 649 | 650 | if xlim is not None: 651 | self.canvas.xlim(xlim) 652 | 653 | if ylim is not None: 654 | self.canvas.ylim(xlim) 655 | 656 | return self.draw() 657 | 658 | def draw(self): 659 | self.output_buffer = [[" "] * self.canvas.y_size for i in range(self.canvas.x_size)] 660 | if self.draw_axes: 661 | self._draw_axes() 662 | 663 | for dk in self.data: 664 | self._plot_data(dk) 665 | 666 | if self.plot_labels: 667 | self._plot_labels() 668 | trans_result = _transpose(_y_reverse(self.output_buffer)) 669 | result = self.new_line.join(["".join(row) for row in trans_result]) 670 | return result 671 | 672 | 673 | def plot(x, y=None, marker=None, shape=(50, 20), draw_axes=True, 674 | newline='\n', plot_slope=False, x_margin=0.05, 675 | y_margin=0.1, plot_labels=True, xlim=None, ylim=None): 676 | 677 | flags = {'shape': shape, 678 | 'draw_axes': draw_axes, 679 | 'newline': newline, 680 | 'marker': marker, 681 | 'plot_slope': plot_slope, 682 | 'margins': (x_margin, y_margin), 683 | 'plot_labels': plot_labels } 684 | 685 | p = AFigure(**flags) 686 | 687 | print p.plot(x, y, marker=marker, plot_slope=plot_slope) 688 | 689 | 690 | def steppify(x, y): 691 | """ Steppify a curve (x,y). Useful for manually filling histograms """ 692 | dx = 0.5 * (x[1:] + x[:-1]) 693 | xx = np.zeros( 2 * len(dx), dtype=float) 694 | yy = np.zeros( 2 * len(y), dtype=float) 695 | xx[0::2], xx[1::2] = dx, dx 696 | yy[0::2], yy[1::2] = y, y 697 | xx = np.concatenate(([x[0] - (dx[0] - x[0])], xx, [x[-1] + (x[-1] - dx[-1])])) 698 | return xx, yy 699 | 700 | 701 | def stemify(x, y): 702 | """ Steppify a curve (x,y). Useful for manually filling histograms """ 703 | xx = np.zeros( 3 * len(x), dtype=float) 704 | yy = np.zeros( 3 * len(y), dtype=float) 705 | xx[0::3], xx[1::3], xx[2::3] = x, x, x 706 | yy[1::3] = y 707 | return xx, yy 708 | 709 | 710 | def hist(x, bins=10, normed=False, weights=None, density=None, histtype='stem', 711 | shape=(50, 20), draw_axes=True, newline='\n', 712 | marker='_.', plot_slope=False, x_margin=0.05, 713 | y_margin=0.1, plot_labels=True, xlim=None, ylim=None ): 714 | 715 | from numpy import histogram 716 | 717 | if histtype not in ['None', 'stem', 'step']: 718 | raise ValueError('histtype must be in [None, stem, step]') 719 | 720 | n, b = histogram(x, bins=bins, range=xlim, normed=normed, weights=weights, density=density) 721 | 722 | _x = 0.5 * ( b[:-1] + b[1:] ) 723 | if histtype == 'step': 724 | step(_x, n.astype(float)) 725 | elif histtype == 'stem': 726 | stem(_x, n.astype(float)) 727 | else: 728 | _y = n.astype(float) 729 | plot(_x, _y, shape=shape, draw_axes=draw_axes, newline=newline, marker=marker, 730 | plot_slope=plot_slope, x_margin=x_margin, y_margin=y_margin, 731 | plot_labels=plot_labels, xlim=xlim, ylim=ylim) 732 | 733 | 734 | def step(x, y, shape=(50, 20), draw_axes=True, 735 | newline='\n', marker='_.', plot_slope=True, x_margin=0.05, 736 | y_margin=0.1, plot_labels=True, xlim=None, ylim=None ): 737 | 738 | _x, _y = steppify(x, y) 739 | plot(_x, _y, shape=shape, draw_axes=draw_axes, newline=newline, marker=marker, 740 | plot_slope=plot_slope, x_margin=x_margin, y_margin=y_margin, 741 | plot_labels=plot_labels, xlim=xlim, ylim=ylim) 742 | 743 | 744 | def stem(x, y, shape=(50, 20), draw_axes=True, 745 | newline='\n', marker='_.', plot_slope=True, x_margin=0.05, 746 | y_margin=0.1, plot_labels=True, xlim=None, ylim=None ): 747 | 748 | _x, _y = stemify(x, y) 749 | plot(_x, _y, shape=shape, draw_axes=draw_axes, newline=newline, marker=marker, 750 | plot_slope=plot_slope, x_margin=x_margin, y_margin=y_margin, 751 | plot_labels=plot_labels, xlim=xlim, ylim=ylim) 752 | 753 | 754 | def hist2d(x, y, bins=[50, 20], range=None, normed=False, weights=None, ncolors=16, 755 | width=50, percentiles=None): 756 | 757 | im, ex, ey = np.histogram2d(x, y, bins, range=None, normed=normed, weights=weights) 758 | 759 | if percentiles is None: 760 | imshow(im, extent=[min(ex), max(ex), min(ey), max(ey)], 761 | ncolors=ncolors, width=width) 762 | else: 763 | percentile_imshow(im, levels=percentiles, extent=None, 764 | width=width, ncolors=width) 765 | 766 | 767 | def percentile_imshow(im, levels=[68, 95, 99], extent=None, width=50, ncolors=16): 768 | _im = im.astype(float) 769 | _im -= im.min() 770 | _im /= _im.max() 771 | 772 | n = len(levels) 773 | for e, lk in enumerate(sorted(levels)): 774 | _im[ _im <= 0.01 * float(lk) ] = n - e 775 | 776 | imshow(1. - _im, extent=None, width=width, ncolors=ncolors) 777 | 778 | 779 | def imshow(im, extent=None, width=50, ncolors=16): 780 | from scipy import ndimage 781 | 782 | width0 = im.shape[0] 783 | _im = ndimage.zoom(im.astype(float), float(width) / float(width0) ) 784 | 785 | _im -= im.min() 786 | _im /= _im.max() 787 | 788 | width, height = _im.shape[:2] 789 | 790 | if len(im.shape) > 2: 791 | _clr = True 792 | else: 793 | _clr = False 794 | 795 | if ncolors == 16: 796 | color = "MNHQ$OC?7>!:-;. "[::-1] 797 | else: 798 | color = '''$@B%8&WM#*oahkbdpqwmZO0QLCJUYXzcvunxrjft/\|()1{}[]?-_+~<>i!lI;:,"^`'. '''[::-1] 799 | ncolors = len(color) 800 | 801 | string = "" 802 | if not _clr: 803 | for h in xrange(height): # first go through the height, otherwise will roate 804 | for w in xrange(width): 805 | string += color[int(_im[w, h] * (ncolors - 1) )] 806 | string += "\n" 807 | else: 808 | for h in xrange(height): # first go through the height, otherwise will roate 809 | for w in xrange(width): 810 | string += color[int(sum(_im[w, h]) * (ncolors - 1) )] 811 | string += "\n" 812 | print string 813 | -------------------------------------------------------------------------------- /stf/common/colors.py: -------------------------------------------------------------------------------- 1 | # This file is part of Viper - https://github.com/botherder/viper 2 | # See the file 'LICENSE' for copying permission. 3 | 4 | import os 5 | import sys 6 | 7 | def color(text, color_code, readline=False): 8 | """Colorize text. 9 | @param text: text. 10 | @param color_code: color. 11 | @return: colorized text. 12 | """ 13 | # $TERM under Windows: 14 | # cmd.exe -> "" (what would you expect..?) 15 | # cygwin -> "cygwin" (should support colors, but doesn't work somehow) 16 | # mintty -> "xterm" (supports colors) 17 | if sys.platform == "win32" and os.getenv("TERM") != "xterm": 18 | return text 19 | if readline: 20 | # special readline escapes to fix colored input promps 21 | # http://bugs.python.org/issue17337 22 | return "\x01\x1b[%dm\x02%s\x01\x1b[0m\x02" % (color_code, text) 23 | return "\x1b[%dm%s\x1b[0m" % (color_code, text) 24 | 25 | def black(text, readline=False): 26 | return color(text, 30, readline) 27 | 28 | def red(text, readline=False): 29 | return color(text, 31, readline) 30 | 31 | def green(text, readline=False): 32 | return color(text, 32, readline) 33 | 34 | def yellow(text, readline=False): 35 | return color(text, 33, readline) 36 | 37 | def blue(text, readline=False): 38 | return color(text, 34, readline) 39 | 40 | def magenta(text, readline=False): 41 | return color(text, 35, readline) 42 | 43 | def cyan(text, readline=False): 44 | return color(text, 36, readline) 45 | 46 | def white(text, readline=False): 47 | return color(text, 37, readline) 48 | 49 | def bold(text, readline=False): 50 | return color(text, 1, readline) 51 | -------------------------------------------------------------------------------- /stf/common/markov_chains.py: -------------------------------------------------------------------------------- 1 | # This file is from the Stratosphere Testing Framework 2 | # See the file 'LICENSE' for copying permission. 3 | 4 | # Library to compute some markov chain functions for the Stratosphere Project. We created them because pykov lacked the second order markov chains 5 | 6 | import math 7 | import sys 8 | 9 | class Matrix(dict): 10 | """ The basic matrix object """ 11 | def __init__(self, *args, **kw): 12 | super(Matrix,self).__init__(*args, **kw) 13 | self.itemlist = super(Matrix,self).keys() 14 | 15 | def set_init_vector(self, init_vector): 16 | self.init_vector = init_vector 17 | 18 | def get_init_vector(self): 19 | return self.init_vector 20 | 21 | def walk_probability(self, states): 22 | """ Compute the probability of generating these states using ourselves. The returned value must be log. """ 23 | try: 24 | #print '\t\twalk_probability' 25 | #print '\t\tReceived states {}'.format(states) 26 | #print '\t\tself init_vector: {}'.format(self.get_init_vector()) 27 | #print '\t\tself matrix: {}'.format(self) 28 | cum_prob = 0 29 | index = 0 30 | # index should be < that len - 1 because index starts in 0, and a two position vector has len 2, but the index of the last position is 1. 31 | # The len of the states should be > 1 because a state of only one char does NOT have any transition. 32 | while index < len(states) - 1 and len(states) > 1: 33 | statestuple = (states[index], states[index + 1]) 34 | #print '\t\ttuple to search: {}'.format(statestuple) 35 | try: 36 | prob12 = math.log(float(self[statestuple])) 37 | #print '\t\tValue for this tuple: {}'.format(self[statestuple]) 38 | #print '\t\tprob12 inside {} (decimal {})'.format(prob12, math.exp(prob12)) 39 | except KeyError: 40 | # The transition is not in the matrix 41 | #print '\t\twalk key error. The transition is not in the matrix' 42 | #prob12 = float('-inf') 43 | cum_prob = float('-inf') 44 | break 45 | #except IndexError: 46 | #print '\t\twalk index error' 47 | cum_prob += prob12 48 | #print '\t\ttotal prob so far {}'.format(cum_prob) 49 | index += 1 50 | #print '\t\tFinal Prob (log): {}'.format(cum_prob) 51 | return cum_prob 52 | except Exception as err: 53 | print type(err) 54 | print err.args 55 | print err 56 | sys.exit(-1) 57 | 58 | 59 | def maximum_likelihood_probabilities(states, order=1): 60 | """ Our own second order Markov Chain implementation """ 61 | initial_matrix = {} 62 | initial_vector = {} 63 | total_transitions = 0 64 | amount_of_states = len(states) 65 | #print 'Receiving {} states to compute the Markov Matrix of {} order'.format(amount_of_states, order) 66 | # 1st order 67 | if order == 1: 68 | # Create matrix 69 | index = 0 70 | while index < amount_of_states: 71 | state1 = states[index] 72 | try: 73 | state2 = states[index + 1] 74 | except IndexError: 75 | # The last state is alone. There is no transaction, forget about it. 76 | break 77 | try: 78 | temp = initial_matrix[state1] 79 | except KeyError: 80 | # First time there is a transition FROM state1 81 | initial_matrix[state1] = {} 82 | initial_vector[state1] = 0 83 | try: 84 | value = initial_matrix[state1][state2] 85 | initial_matrix[state1][state2] = value + 1 86 | except KeyError: 87 | # First time there is a transition FROM state 1 to state2 88 | initial_matrix[state1][state2] = 1 89 | initial_vector[state1] += 1 90 | total_transitions += 1 91 | # Move along 92 | index += 1 93 | # Normalize using the initial vector 94 | matrix = Matrix() 95 | init_vector = {} 96 | for state1 in initial_matrix: 97 | # Create the init vector 98 | init_vector[state1] = initial_vector[state1] / float(total_transitions) 99 | for state2 in initial_matrix[state1]: 100 | value = initial_matrix[state1][state2] 101 | initial_matrix[state1][state2] = value / float(initial_vector[state1]) 102 | # Change the style of the matrix 103 | matrix[(state1,state2)] = initial_matrix[state1][state2] 104 | matrix.set_init_vector(init_vector) 105 | #print init_vector 106 | #for value in matrix: 107 | # print value, matrix[value] 108 | return (init_vector, matrix) 109 | 110 | 111 | -------------------------------------------------------------------------------- /stf/common/out.py: -------------------------------------------------------------------------------- 1 | # This file was mostly taken from Viper - https://github.com/botherder/viper 2 | # See the file 'LICENSE' for copying permission. 3 | 4 | from prettytable import PrettyTable 5 | 6 | from stf.common.colors import * 7 | 8 | def print_info(message): 9 | print(bold(cyan("[*]")) + " {0}".format(message)) 10 | 11 | def print_item(message, tabs=0): 12 | print(" {0}".format(" " * tabs) + cyan("-") + " {0}".format(message)) 13 | 14 | def print_row(data): 15 | """ Intended for long tables. We want to see the output quickly and not wait some minutes until the table is created """ 16 | print('| '), 17 | for datum in data: 18 | print('{:80}'.format(datum)), 19 | print('| '), 20 | print 21 | 22 | def print_warning(message): 23 | print(bold(yellow("[!]")) + " {0}".format(message)) 24 | 25 | def print_error(message): 26 | print(bold(red("[!]")) + " {0}".format(message)) 27 | 28 | def print_success(message): 29 | print(bold(green("[+]")) + " {0}".format(message)) 30 | 31 | def table(header, rows): 32 | table = PrettyTable(header) 33 | table.align = 'l' 34 | table.padding_width = 1 35 | 36 | for row in rows: 37 | table.add_row(row) 38 | 39 | return table 40 | -------------------------------------------------------------------------------- /stf/core/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/stratosphereips/StratosphereTestingFramework/381eca74358dfa3847727314f2ae615cd5c8170a/stf/core/__init__.py -------------------------------------------------------------------------------- /stf/core/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/stratosphereips/StratosphereTestingFramework/381eca74358dfa3847727314f2ae615cd5c8170a/stf/core/__init__.pyc -------------------------------------------------------------------------------- /stf/core/configuration.py: -------------------------------------------------------------------------------- 1 | import ConfigParser 2 | import os 3 | 4 | from stf.common.out import * 5 | 6 | 7 | class Configuration(object): 8 | """ 9 | The class to deal with the configuration. Every other class that wants to read the configuration just instantiate this one. 10 | """ 11 | def __init__(self): 12 | self.config = ConfigParser.ConfigParser() 13 | self.zeoconffile = None 14 | self.zodbconffile = None 15 | 16 | def read_conf_file(self, conf_file): 17 | """ Read the conf file""" 18 | # Read the sections 19 | self.config.read(conf_file) 20 | stf_section = self.ConfigSectionMap('stf') 21 | try: 22 | self.zeoconffile = stf_section['zeoconfigurationfile'] 23 | self.zodbconffile = stf_section['zodbconfigurationfile'] 24 | except: 25 | print_error('Errors in the configuration files.') 26 | return False 27 | return True 28 | 29 | def get_zeoconf_file(self): 30 | return self.zeoconffile 31 | 32 | def get_zodbconf_file(self): 33 | return self.zodbconffile 34 | 35 | def ConfigSectionMap(self,section): 36 | """ Taken from the python site""" 37 | dict1 = {} 38 | options = self.config.options(section) 39 | for option in options: 40 | try: 41 | dict1[option] = self.config.get(section, option) 42 | if dict1[option] == -1: 43 | print_info("Skip: %s" % option) 44 | except: 45 | print_error("Exception on %s!" % option) 46 | dict1[option] = None 47 | return dict1 48 | 49 | 50 | __configuration__ = Configuration() 51 | -------------------------------------------------------------------------------- /stf/core/database.py: -------------------------------------------------------------------------------- 1 | import ZODB.config 2 | import persistent 3 | import transaction 4 | import time 5 | from datetime import datetime 6 | 7 | from stf.core.configuration import __configuration__ 8 | from stf.common.out import * 9 | from stf.core.dataset import __datasets__ 10 | from stf.core.connections import __group_of_group_of_connections__ 11 | from stf.core.models import __groupofgroupofmodels__ 12 | from stf.core.notes import __notes__ 13 | from stf.core.labels import __group_of_labels__ 14 | 15 | class Database: 16 | def __init__(self): 17 | """ Initialize """ 18 | pass 19 | 20 | def start(self): 21 | """ From some reason we should initialize the db from a method, we can not do it in the constructor """ 22 | dbconffile = __configuration__.get_zodbconf_file() 23 | self.db = ZODB.config.databaseFromURL(dbconffile) 24 | 25 | # The server and port should be read from a future configuration 26 | self.connection = self.db.open() 27 | self.root = self.connection.root() 28 | 29 | 30 | # Datasets 31 | try: 32 | __datasets__.datasets = self.root['datasets'] 33 | except KeyError: 34 | self.root['datasets'] = __datasets__.datasets 35 | 36 | # Connections 37 | try: 38 | __group_of_group_of_connections__.group_of_connections = self.root['connections'] 39 | except KeyError: 40 | self.root['connections'] = __group_of_group_of_connections__.group_of_connections 41 | 42 | # Models 43 | try: 44 | __groupofgroupofmodels__.group_of_models = self.root['models'] 45 | except KeyError: 46 | self.root['models'] = __groupofgroupofmodels__.group_of_models 47 | 48 | # Notes 49 | try: 50 | __notes__.notes = self.root['notes'] 51 | except KeyError: 52 | self.root['notes'] = __notes__.notes 53 | 54 | # Labels 55 | try: 56 | __group_of_labels__.labels = self.root['labels'] 57 | except KeyError: 58 | self.root['labels'] = __group_of_labels__.labels 59 | 60 | 61 | def has_structure(self, structure_name): 62 | """ This method searches for a structure in the db""" 63 | for structure in self.root: 64 | if structure == structure_name: 65 | return True 66 | return False 67 | 68 | def get_new_structure(self, structure): 69 | """ Given a structure, set the main dict from the db """ 70 | name = str(structure.get_name()) 71 | return self.root[name] 72 | 73 | def get_structures(self): 74 | """ get all the structures """ 75 | return self.root 76 | 77 | def set_new_structure(self, structure): 78 | """ 79 | This method takes an object from a new structure (typically from a module) and keeps record of it in the database. 80 | A strcture is the main object from the module that we want to store in the db. Actually we store its main dict. 81 | """ 82 | try: 83 | name = structure.get_name() 84 | except AttributeError: 85 | print_error('The new registered structure does not implement get_name()') 86 | return False 87 | try: 88 | main_dict = structure.get_main_dict() 89 | print_info('Registering structure name: {}'.format(name)) 90 | self.root[name] = main_dict 91 | return True 92 | except AttributeError: 93 | print_error('The structure does not implement get_main_dict()') 94 | return False 95 | 96 | def list(self): 97 | #for structure in self.root: 98 | # print_info('Amount of {} in the DB so far: {}'.format(structure, len(self.root[structure]))) 99 | pass 100 | 101 | def delete_structure(self, structure_name): 102 | """ Delete a structure from the db """ 103 | try: 104 | structure = self.root[structure_name] 105 | print_warning('Are you sure you want to delete the structure {} from the db? (YES/NO)'.format(structure_name)) 106 | input = raw_input() 107 | if input == "YES": 108 | self.root.pop(structure_name) 109 | print_info('Structure {} deleted from the db'.format(structure_name)) 110 | except KeyError: 111 | print_error('No Structure name available.') 112 | 113 | def list_structures(self): 114 | for structure in self.root: 115 | print_info('Structure: {}. Amount of objects in db: {}'.format(structure, len(self.root[structure]))) 116 | 117 | def close(self): 118 | """ Close the db """ 119 | transaction.commit() 120 | self.connection.close() 121 | # In the future we should try to pack based on time 122 | self.db.pack 123 | self.db.close() 124 | 125 | def info(self): 126 | """ Info about the database""" 127 | print_warning('Info about the root object of the database') 128 | print_info('Main Branch | Len') 129 | for mainbranches in self.root.keys(): 130 | print('\t{:12} | {}'.format(mainbranches, len(self.root[mainbranches]))) 131 | print_info('_p_changed (the persistent state of the object): {}'.format(self.root._p_changed)) 132 | print('\t-> None: The object is a ghost.\n\t-> False but not None: The object is saved (or has never been saved).\n\t-> True: The object has been modified since it was last saved.') 133 | print_info('_p_state (object persistent state token): {}'.format('GHOST' if self.root._p_state == -1 else 'UPTODATE' if self.root._p_state == 0 else 'CHANGED' if self.root._p_state == 1 else 'STICKY')) 134 | print_info('_p_jar: {}'.format(self.root._p_jar)) 135 | print_info('_p_oid (persistent object id): {}'.format(self.root._p_oid)) 136 | print 137 | print_warning('Info about the database object ') 138 | print_info('Database Name: {}.'.format(self.db.getName())) 139 | print_info('Database Size: {} B ({} MB).'.format(self.db.getSize(), self.db.getSize()/1024/1024)) 140 | print_info('Object Count: {}.'.format(self.db.objectCount())) 141 | print_info('Connection debug info: {}.'.format(self.db.connectionDebugInfo())) 142 | print_info('Cache details:') 143 | for detail in self.db.cacheDetail(): 144 | print_info('\t{}'.format(detail)) 145 | 146 | def revert(self): 147 | """ revert the connection of the database to the previous state before the last pack""" 148 | question=raw_input('Warning, this command sync the database on disk with the information on memory. Effectively erasing the changes that were not committed and bringing the new information committed by other instances. \n YES or NO?:') 149 | if question == 'YES': 150 | self.connection.sync() 151 | else: 152 | print_error('Not reverting.') 153 | 154 | def pack(self): 155 | """ Pack the database """ 156 | self.db.pack() 157 | 158 | def commit(self): 159 | """ Commit the changes in the connection to the database """ 160 | import transaction 161 | transaction.commit() 162 | 163 | 164 | 165 | 166 | 167 | __database__ = Database() 168 | -------------------------------------------------------------------------------- /stf/core/dataset.py: -------------------------------------------------------------------------------- 1 | # This file was partially taken from Viper 2 | # See the file 'LICENSE' for copying permission. 3 | 4 | import time 5 | import datetime 6 | import persistent 7 | import BTrees.IOBTree 8 | import transaction 9 | import os 10 | from subprocess import Popen,PIPE 11 | 12 | from stf.common.out import * 13 | from stf.core.file import File 14 | from stf.core.notes import __notes__ 15 | 16 | 17 | ########################### 18 | ########################### 19 | ########################### 20 | class Dataset(persistent.Persistent): 21 | """ 22 | The Dataset class. 23 | """ 24 | def __init__(self, id): 25 | self.id = id 26 | self.name = None 27 | # Timestamp of the creation of the session. 28 | self.added_on = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d %H:%M:%S') 29 | self.ctime = None 30 | # Dict of files related to this dataset 31 | self.files = {} 32 | # Foler that holds all the files for this dataset 33 | self.folder = None 34 | # This dict holds all the groups of models related with this dataset. There can be many because there are several models constructors. Is a dict only to search faster. 35 | self.group_of_models = {} 36 | # This stores the id of the group of connections related with this dataset. Only one group of connections is related. 37 | self.group_of_connections_id = False 38 | 39 | def get_file_type(self,type): 40 | """ Return the file with type x in this dataset""" 41 | for file in self.get_files(): 42 | if file.get_type() == type: 43 | return file 44 | return False 45 | 46 | def get_files(self): 47 | """ Return the vector of files of the dataset""" 48 | return self.files.values() 49 | 50 | def get_file(self, id): 51 | """ Return a file object given its id """ 52 | try: 53 | return self.files[int(id)] 54 | except KeyError: 55 | print_error('No such file id') 56 | 57 | def get_id(self): 58 | return self.id 59 | 60 | def get_name(self): 61 | return self.name 62 | 63 | def get_atime(self): 64 | return self.added_on 65 | 66 | def is_current(self): 67 | return self.is_current 68 | 69 | def set_name(self,name): 70 | self.name = name 71 | 72 | def get_main_file(self): 73 | """ Returns the name of the first file used to create the dataset. Usually the only one""" 74 | try: 75 | for file in self.files: 76 | return self.files[file] 77 | except KeyError: 78 | print_error('There is no main file in this dataset!') 79 | return False 80 | 81 | def del_file(self,fileid): 82 | """ Delete a file from the dataset""" 83 | try: 84 | print_info('File {} with id {} deleted from dataset {}'.format(self.files[fileid].get_name(), self.files[fileid].get_id(), self.get_name() )) 85 | self.files.pop(fileid) 86 | # If this was the last file in the dataset, delete the dataset 87 | if len(self.files) == 0: 88 | __datasets__.delete(__datasets__.current.get_id()) 89 | except KeyError: 90 | print_error('No such file id.') 91 | 92 | 93 | def add_file(self,filename): 94 | """ Add a file to this dataset. """ 95 | # Check that the file exists 96 | if not os.path.exists(filename) or not os.path.isfile(filename): 97 | print_error('File not found: {}'.format(filename)) 98 | return 99 | # We should have only one file per type 100 | short_name = os.path.split(filename)[1] 101 | extension = short_name.split('.')[-1] 102 | if 'xz' in extension: 103 | # The file is compressed, but argus can deal with it. 104 | if 'biargus' in short_name.split('.')[-2]: 105 | extension = 'biargus' 106 | # search for the extensions of the files in the dataset 107 | for file in self.files: 108 | if extension in self.files[file].get_type(): 109 | print_error('Only one type of file per dataset is allowed.') 110 | return False 111 | # Get the new id for this file 112 | try: 113 | # Get the id of the last file in the dataset 114 | f_id = self.files[list(self.files.keys())[-1]].get_id() + 1 115 | except (KeyError, IndexError): 116 | f_id = 0 117 | # Create the file object 118 | f = File(filename, f_id) 119 | # Add it to the list of files related to this dataset 120 | self.files[f_id] = f 121 | print_info('Added file {} to dataset {}'.format(filename, self.name)) 122 | 123 | def get_folder(self): 124 | return self.folder 125 | 126 | def set_folder(self, folder): 127 | self.folder = folder 128 | 129 | def list_files(self): 130 | rows = [] 131 | for file in self.files.values(): 132 | rows.append([file.get_short_name(), file.get_id() , file.get_modificationtime(), file.get_size_in_megabytes(), file.get_type()]) 133 | print(table(header=['File Name', 'Id', 'Creation Time', 'Size', 'Type'], rows=rows)) 134 | 135 | def info_about_file(self,file_id): 136 | file = self.files[int(file_id)] 137 | file.info() 138 | 139 | def generate_biargus(self): 140 | """ Generate the biargus file from the pcap. We know that there is a pcap in the dataset""" 141 | print_info('Generating the biargus file.') 142 | pcap_file_name = self.get_file_type('pcap').get_name() 143 | pcap_file_name_without_extension = '.'.join(pcap_file_name.split('.')[:-1]) 144 | biargus_file_name = pcap_file_name_without_extension + '.biargus' 145 | try: 146 | argus_path = Popen('bash -i -c "type argus"', shell=True, stdin=PIPE, stdout=PIPE).communicate()[0].split()[0] 147 | except IndexError: 148 | print_error('argus is not installed. We can not generate the flow files. Download and install from http://qosient.com/argus/dev/argus-clients-latest.tar.gz and http://qosient.com/argus/dev/argus-latest.tar.gz') 149 | return False 150 | if argus_path: 151 | # If an .biargus file already exist, we must delete it because argus appends the output 152 | (data, error) = Popen('rm -rf '+biargus_file_name, shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE).communicate() 153 | (argus_data,argus_error) = Popen('argus -F ./confs/argus.conf -r '+pcap_file_name+' -w '+biargus_file_name, shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE).communicate() 154 | if not argus_error: 155 | # Add the new biargus file to the dataset 156 | self.add_file(biargus_file_name) 157 | else: 158 | print_error('There was an error with argus.') 159 | return False 160 | return True 161 | else: 162 | print_error('argus is not installed. We can not generate the flow files. Download and install from http://qosient.com/argus/dev/argus-clients-latest.tar.gz and http://qosient.com/argus/dev/argus-latest.tar.gz') 163 | return False 164 | 165 | def generate_binetflow(self): 166 | """ Generate the binetflow file from the biargus. We know that there is a biargus in the dataset""" 167 | print_info('Generating the binetflow file.') 168 | try: 169 | biargus_file_name = self.get_file_type('biargus').get_name() 170 | except AttributeError: 171 | print_error('Can not generate the biargus file. Maybe we don\'t have permisions in the file or folder?') 172 | return False 173 | biargus_file_name_without_extension = '.'.join(biargus_file_name.split('.')[:-1]) 174 | binetflow_file_name = biargus_file_name_without_extension + '.binetflow' 175 | ra_path = Popen('bash -i -c "type ra"', shell=True, stdin=PIPE, stdout=PIPE).communicate()[0].split()[0] 176 | if ra_path: 177 | (ra_data,ra_error) = Popen('ra -F ./confs/ra.conf -n -Z b -r '+biargus_file_name+'|sort -n > '+binetflow_file_name, shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE).communicate() 178 | if not ra_error: 179 | # Add the new biargus file to the dataset 180 | self.add_file(binetflow_file_name) 181 | else: 182 | print_error('There was an error with ra.') 183 | print ra_error 184 | return False 185 | return True 186 | else: 187 | print_error('ra is not installed. We can not generate the flow files. Download and install from http://qosient.com/argus/dev/argus-clients-latest.tar.gz and http://qosient.com/argus/dev/argus-latest.tar.gz') 188 | return False 189 | 190 | 191 | def __repr__(self): 192 | return (' > Dataset id {}, and name {}.'.format(self.id, self.name)) 193 | 194 | def add_group_of_models(self, group_of_models_id): 195 | """ Add the group of models id to the list of groups of models related with this dataset """ 196 | self.group_of_models[group_of_models_id] = None 197 | 198 | def get_group_of_models(self): 199 | return self.group_of_models 200 | 201 | def has_group_of_models(self,group_of_models_id): 202 | return self.group_of_models.has_key(group_of_models_id) 203 | 204 | def add_group_of_connections_id(self, group_of_connections_id): 205 | """ Add the group of connections that is related to this dataset. Is only one, but we store it this way. """ 206 | self.group_of_connections_id = group_of_connections_id 207 | 208 | def remove_group_of_connections_id(self, group_of_connections_id): 209 | self.group_of_connections_id = False 210 | 211 | def get_group_of_connections_id(self): 212 | try: 213 | return self.group_of_connections_id 214 | except AttributeError: 215 | return False 216 | 217 | def set_group_of_connections_id(self, group_of_connections_id): 218 | self.group_of_connections_id = group_of_connections_id 219 | 220 | def set_note_id(self, note_id): 221 | self.note_id = note_id 222 | 223 | def edit_note(self): 224 | """ Edit the note related with this dataset or create a new one and edit it """ 225 | try: 226 | note_id = self.note_id 227 | __notes__.edit_note(note_id) 228 | except AttributeError: 229 | self.note_id = __notes__.new_note() 230 | __notes__.edit_note(self.note_id) 231 | 232 | def edit_folder(self, new_folder): 233 | """ Edit the folder related with this dataset """ 234 | # First change it in the dataset 235 | self.set_folder(new_folder) 236 | # Now change it in the files inside the dataset 237 | for file in self.get_files(): 238 | filename = file.get_name() 239 | real_file_name = os.path.split(filename)[1] 240 | # does the folder has a final / ? 241 | current_folder = self.get_folder() 242 | if current_folder[-1] != '/': 243 | current_folder += '/' 244 | newfilenanme = current_folder + real_file_name 245 | file.set_name(newfilenanme) 246 | 247 | 248 | def del_note(self): 249 | """ Delete the note related with this dataset """ 250 | try: 251 | # First delete the note 252 | note_id = self.note_id 253 | __notes__.del_note(note_id) 254 | # Then delete the reference to the note 255 | del self.note_id 256 | 257 | except AttributeError: 258 | print_error('No such note id exists.') 259 | 260 | def get_note_id(self): 261 | """ Return the note id or false """ 262 | try: 263 | return self.note_id 264 | except AttributeError: 265 | return False 266 | 267 | 268 | 269 | 270 | ########################### 271 | ########################### 272 | ########################### 273 | class Datasets(persistent.Persistent): 274 | def __init__(self): 275 | self.current = False 276 | # The main dictionary of datasets objects using its id as index 277 | self.datasets = BTrees.IOBTree.BTree() 278 | 279 | def get_datasets_ids(self): 280 | """ Return the ids of the datasets """ 281 | return self.datasets.keys() 282 | 283 | def get_dataset(self, id): 284 | """ Return a dataset objet given the id """ 285 | try: 286 | return self.datasets[id] 287 | except: 288 | print_error('No such dataset id') 289 | return False 290 | 291 | def delete(self, dataset_id): 292 | """ Delete a dataset from the list of datasets """ 293 | # Verify the decision 294 | print_warning('Warning! The connections and models created from this dataset WILL be deleted. However, all the labels created from this dataset WILL NOT be deleted (due to their associations with third-party modules). You should delete the labels by hand.') 295 | input = raw_input('Are you sure you want to delete this dataset? (YES/NO): ') 296 | if input != 'YES': 297 | return False 298 | # Before deleting the dataset, delete the connections 299 | from stf.core.connections import __group_of_group_of_connections__ 300 | __group_of_group_of_connections__.delete_group_of_connections(dataset_id) 301 | # Before deleting the dataset, delete the models 302 | from stf.core.models import __groupofgroupofmodels__ 303 | __groupofgroupofmodels__.delete_group_of_models_with_dataset_id(dataset_id) 304 | 305 | try: 306 | # Now delete the dataset 307 | self.datasets.pop(dataset_id) 308 | print_info("Deleted dataset #{0}".format(dataset_id)) 309 | # If it was the current dataset, we have no current 310 | if self.current and self.current.get_id() == dataset_id: 311 | self.current = False 312 | except ValueError: 313 | print_info('You should give an dataset id') 314 | except KeyError: 315 | print_info('Dataset ID non existant.') 316 | 317 | def add_file(self, filename): 318 | """ Add a new file to the current dataset""" 319 | if self.current: 320 | # Check that the file exists 321 | if not os.path.exists(filename) or not os.path.isfile(filename): 322 | print_error('File not found: {}'.format(filename)) 323 | return 324 | # Add this file to the dataset 325 | self.current.add_file(filename) 326 | else: 327 | print_error('No dataset selected. Use -s option.') 328 | 329 | def del_file(self,fileid): 330 | """ Delete a file to the current dataset""" 331 | if self.current: 332 | # Delete this file from the dataset 333 | self.current.del_file(int(fileid)) 334 | 335 | else: 336 | print_error('No dataset selected. Use -s option.') 337 | 338 | def create(self,filename): 339 | """ Create a new dataset from a file name""" 340 | 341 | # Check that the file exists 342 | if not os.path.exists(filename) or not os.path.isfile(filename): 343 | print_error('File not found: {}'.format(filename)) 344 | return 345 | 346 | # Get the new id for this dataset 347 | try: 348 | # Get the id of the last dataset in the database 349 | dat_id = self.datasets[list(self.datasets.keys())[-1]].get_id() + 1 350 | except (KeyError, IndexError): 351 | dat_id = 0 352 | 353 | # Create the dataset object 354 | dataset = Dataset(dat_id) 355 | 356 | # Ask for a dataset name or default to the file name 357 | name = raw_input('The name of the dataset or \'Enter\' to use the name of the last folder:') 358 | if not name: 359 | name = os.path.split(filename)[0].split('/')[-1] 360 | 361 | # Set the name 362 | dataset.set_name(name) 363 | 364 | # Add this file to the dataset 365 | dataset.add_file(filename) 366 | 367 | # Store the folder of this dataset 368 | folder = os.path.split(filename)[0] 369 | dataset.set_folder(folder) 370 | 371 | # Add th enew dataset to the dict 372 | self.datasets[dataset.get_id()] = dataset 373 | print_info("Dataset {} added with id {}.".format(name, dataset.get_id())) 374 | self.current = dataset 375 | 376 | def list(self): 377 | """ List all the datasets """ 378 | print_info("Datasets Available:") 379 | rows = [] 380 | for dataset in self.datasets.values(): 381 | main_file = dataset.get_main_file() 382 | rows.append([dataset.get_name(), dataset.get_id() , dataset.get_atime() , main_file.get_short_name(), main_file.get_modificationtime(), dataset.get_folder(), True if (self.current and self.current.get_id() == dataset.get_id()) else False, dataset.get_note_id() if dataset.get_note_id() >= 0 else '' ]) 383 | print(table(header=['Dataset Name', 'Id', 'Added Time', 'Main File Name', 'Main File Creation Time', 'Folder', 'Current', 'Note'], rows=rows)) 384 | 385 | def list_files(self): 386 | """ List all the files in dataset """ 387 | if self.current: 388 | print_info('Getting information about the files... please wait') 389 | print_info('Files Available in Dataset {}:'.format(self.current.get_name())) 390 | self.current.list_files() 391 | else: 392 | print_error('No dataset selected. Use -s option.') 393 | 394 | def info_about_file(self,file_id): 395 | """ Give info about a specific file in a dataset""" 396 | if self.current: 397 | try: 398 | self.current.info_about_file(int(file_id)) 399 | except (KeyError, UnboundLocalError): 400 | print_info('No such file id') 401 | else: 402 | print_error('No dataset selected. Use -s option.') 403 | 404 | def length(self): 405 | """ Return the length of the dict """ 406 | return len(self.datasets) 407 | 408 | def is_set(self): 409 | """ Does the dict exists?""" 410 | if self.datasets: 411 | return True 412 | else: 413 | return False 414 | 415 | def unselect_current(self): 416 | """ UnSelects the current dataset""" 417 | self.current = False 418 | 419 | def select(self,dataset_id): 420 | """ Selects a dataset as current to enable other more specific commands""" 421 | try: 422 | self.current = self.datasets[int(dataset_id)] 423 | print_info('The current dataset is {} with id {}'.format(self.current.get_name(), self.current.get_id())) 424 | except (KeyError, ValueError): 425 | print_error('No such dataset id') 426 | 427 | def generate_argus_files(self): 428 | """ Generate the biargus and binetflow files""" 429 | if self.current: 430 | # Do we have a binetflow file in the dataset? 431 | binetflow_in_dataset = self.current.get_file_type('binetflow') 432 | # Do we have a binetflow file in the folder? 433 | # Do we have a biargus file in the dataset? 434 | biargus_in_dataset = self.current.get_file_type('biargus') 435 | # Do we have a biargus file in the folder? 436 | 437 | # Do we have a pcap file in the dataset? 438 | pcap_in_dataset = self.current.get_file_type('pcap') 439 | 440 | if binetflow_in_dataset: 441 | # Ask if we should regenerate 442 | pass 443 | elif biargus_in_dataset: 444 | # We should generate the binetflow 445 | # Or regenerate 446 | self.current.generate_binetflow() 447 | elif pcap_in_dataset: 448 | # We should generate the biargus and the binetflow 449 | self.current.generate_biargus() 450 | self.current.generate_binetflow() 451 | else: 452 | print_error('At least a pcap file should be in the dataset.') 453 | 454 | # Do we have a pcap file in the folder? 455 | else: 456 | print_error('No dataset selected. Use -s option.') 457 | 458 | def edit_folder(self, folder_name): 459 | """ Get a dataset id and edit its folder """ 460 | if self.current: 461 | self.current.edit_folder(folder_name) 462 | else: 463 | print_error('No dataset selected. Use -s option.') 464 | 465 | def edit_note(self, dataset_id): 466 | """ Get a dataset id and edit its note """ 467 | try: 468 | dataset = self.datasets[int(dataset_id)] 469 | dataset.edit_note() 470 | except KeyError: 471 | print_error('No such dataset id.') 472 | 473 | def del_note(self, dataset_id): 474 | """ Get a dataset id and delete its note """ 475 | try: 476 | dataset = self.datasets[int(dataset_id)] 477 | dataset.del_note() 478 | except KeyError: 479 | print_error('No such dataset id.') 480 | 481 | 482 | __datasets__ = Datasets() 483 | -------------------------------------------------------------------------------- /stf/core/file.py: -------------------------------------------------------------------------------- 1 | import time 2 | import os 3 | from subprocess import Popen,PIPE 4 | import datetime 5 | from dateutil import parser 6 | import persistent 7 | 8 | from stf.common.out import * 9 | 10 | class File(persistent.Persistent): 11 | """ A Class to store all the info of different file types""" 12 | def __init__(self, filename, id): 13 | # initialize some of the attributes 14 | self.filename = filename 15 | self.id = id 16 | self.duration = False 17 | 18 | # Initialize some inner attributes 19 | self.capinfo = False 20 | self.histoinfo = False 21 | self.binetflowinfo = False 22 | 23 | # Set the modification time 24 | self.set_modificationtime() 25 | # Guess the file type to compute some of the info 26 | self.guess_type() 27 | 28 | def get_size_in_megabytes(self): 29 | size = self.get_size() 30 | if size: 31 | return str(size / 1024.0 / 1024.0)+' MB' 32 | # We can't even compute the size 33 | return "No data" 34 | 35 | def get_size(self): 36 | try: 37 | size = self.size 38 | except (AttributeError, TypeError): 39 | size = self.compute_size() 40 | if size: 41 | self.size = size 42 | else: 43 | # Could not be computed 44 | return False 45 | return self.size 46 | 47 | def set_duration(self,duration): 48 | self.duration = duration 49 | 50 | def get_duration(self): 51 | if not self.duration: 52 | if self.type == 'pcap': 53 | self.get_capinfos() 54 | elif self.type == 'binetflow': 55 | self.get_binetflowinfos() 56 | # To avoid returing 'False' for the duration 57 | if self.duration: 58 | return self.duration 59 | else: 60 | return '' 61 | 62 | def compute_size(self): 63 | try: 64 | size = os.path.getsize(self.get_name()) 65 | except OSError: 66 | print_error('The file is not available in your disk.') 67 | return False 68 | return size 69 | 70 | def get_id(self): 71 | return self.id 72 | 73 | def set_modificationtime(self): 74 | ctime = time.ctime(os.path.getmtime(self.get_name())) 75 | self.ctime = ctime 76 | 77 | def get_modificationtime(self): 78 | return self.ctime 79 | 80 | def get_name(self): 81 | return self.filename 82 | 83 | def get_short_name(self): 84 | """ Only the name of the file without the path""" 85 | return os.path.split(self.filename)[1] 86 | 87 | def set_name(self, filename): 88 | self.filename = filename 89 | 90 | def guess_type(self): 91 | short_name = os.path.split(self.filename)[1] 92 | extension = short_name.split('.')[-1] 93 | if 'xz' in extension: 94 | # The file is compressed, but argus can deal with it. 95 | if 'biargus' in short_name.split('.')[-2]: 96 | extension = 'biargus' 97 | if 'pcap' in extension: 98 | self.set_type('pcap') 99 | elif 'netflow' in extension: 100 | self.set_type('binetflow') 101 | elif 'weblog' in extension: 102 | self.set_type('weblog') 103 | elif 'argus' in extension: 104 | self.set_type('biargus') 105 | elif 'exe' in extension: 106 | self.set_type('exe') 107 | else: 108 | self.set_type('Unknown') 109 | 110 | def set_type(self, type): 111 | self.type = type 112 | 113 | def get_type(self): 114 | """ Returns the type of file """ 115 | return self.type 116 | 117 | def get_binetflowinfos(self): 118 | """ Get info about binetflow files""" 119 | if self.binetflowinfo == False and self.get_type() == 'binetflow': 120 | # Get the time in the first line, ignoring the header 121 | binetflow_first_flow = Popen('head -n 2 '+self.get_name()+'|tail -n 1', shell=True, stdin=PIPE, stdout=PIPE).communicate()[0] 122 | first_flow_date = parser.parse(binetflow_first_flow.split(',')[0]) 123 | 124 | # Get the time in the last line 125 | binetflow_last_flow = Popen('tail -n 1 '+self.get_name(), shell=True, stdin=PIPE, stdout=PIPE).communicate()[0] 126 | last_flow_date = parser.parse(binetflow_last_flow.split(',')[0]) 127 | 128 | # Compute the difference 129 | time_diff = last_flow_date - first_flow_date 130 | self.set_duration(time_diff) 131 | 132 | # Now fill the data for binetflows 133 | self.binetflowinfo = {} 134 | # Duration 135 | self.binetflowinfo['Duration'] = self.get_duration() 136 | 137 | # Amount of flows 138 | amount_of_flows = Popen('wc -l '+self.get_name(), shell=True, stdin=PIPE, stdout=PIPE).communicate()[0].split()[0] 139 | self.binetflowinfo['Amount of flows'] = amount_of_flows 140 | 141 | 142 | # Always return true 143 | return True 144 | 145 | def get_capinfos(self): 146 | """ Get info with capinfos""" 147 | if self.capinfo == False and self.get_type() == 'pcap': 148 | # I don't know how to do this better... 149 | capinfos_path = Popen('bash -i -c "type capinfos"', shell=True, stdin=PIPE, stdout=PIPE).communicate()[0].split()[0] 150 | 151 | if capinfos_path: 152 | (capinfos_data,capinfos_error) = Popen('tcpdump -n -s0 -r ' + self.get_name() + ' -w - | capinfos -r -', shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE).communicate() 153 | capinfos_info = capinfos_data.strip().split('\n') 154 | 155 | # Process the capinfo info 156 | if 'Value too large' in capinfos_error: 157 | print_error('There was an error with capinfos. Maybe the file is too large. See https://www.wireshark.org/lists/wireshark-users/200908/msg00008.html') 158 | self.capinfo = '' 159 | return False 160 | elif capinfos_info: 161 | self.capinfo = {} 162 | for line in capinfos_info: 163 | header = line.split(': ')[0].strip() 164 | info = line.split(': ')[1].strip() 165 | if 'Number of packets' in header: 166 | self.capinfo['Number of packets'] = info 167 | elif 'Capture duration' in header: 168 | self.capinfo['Capture duration'] = datetime.timedelta(seconds=int(info.split()[0])) 169 | # The default duration of the file can be setted now also 170 | self.set_duration(self.capinfo['Capture duration']) 171 | elif 'Start time' in header: 172 | self.capinfo['Start time'] = parser.parse(info) 173 | elif 'End time' in header: 174 | self.capinfo['End time'] = parser.parse(info) 175 | elif 'MD5' in header: 176 | self.capinfo['MD5'] = info 177 | elif 'SHA1' in header: 178 | self.capinfo['SHA1'] = info 179 | 180 | return True 181 | else: 182 | print_error('capinfos is not installed. We can not get more information about the pcap file. apt-get install wireshark-common') 183 | return False 184 | 185 | def get_bytes_histo(self): 186 | """ Use tshark to get the amount of bytes per 10minutes in the pcap file""" 187 | if self.histoinfo == False and self.get_type() == 'pcap': 188 | capinfos_path = Popen('bash -i -c "type tshark"', shell=True, stdin=PIPE, stdout=PIPE).communicate()[0].split()[0] 189 | 190 | if capinfos_path: 191 | (tshark_data,tshark_error) = Popen('tshark -r '+self.get_name()+' -z io,stat,300,"COUNT(frame)frame" -q|grep "<>"|head -n 24', shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE).communicate() 192 | tshark_info = tshark_data.strip().split('\n') 193 | self.histoinfo = {} 194 | for line in tshark_info: 195 | header = line.split('|')[1].strip() 196 | # Store the key as int for later sorting 197 | number_in_header = float(header.split('<>')[0].strip()) 198 | info = line.split('|')[2].strip() 199 | self.histoinfo[number_in_header] = header + '|' + info 200 | return True 201 | else: 202 | print_error('tshark is not installed. We can not get more information about the pcap file. apt-get install tshark') 203 | return False 204 | elif self.histoinfo and self.get_type() == 'pcap': 205 | return True 206 | 207 | def get_md5(self): 208 | """ Return the md5 of the file """ 209 | try: 210 | return self.md5 211 | except AttributeError: 212 | self.md5 = tshark_data = Popen('md5sum '+self.get_name()+' | awk \'{print $1}\'', shell=True, stdin=PIPE, stdout=PIPE).communicate()[0] 213 | return self.md5 214 | 215 | def info(self): 216 | rows = [] 217 | print_info('Information of file name {} with id {}'.format(self.get_short_name(), self.get_id())) 218 | 219 | rows.append(['Type', self.get_type()]) 220 | rows.append(['Creation Time', self.get_modificationtime()]) 221 | 222 | # Get the file size 223 | #if not self.get_size(): 224 | #self.compute_size() 225 | size = self.get_size() 226 | if not size: 227 | return False 228 | rows.append(['Size', str(size/1024.0/1024)+' MB']) 229 | 230 | # Get more info from pcap files 231 | if 'pcap' in self.get_type(): 232 | # Get capinfo data 233 | if self.get_capinfos(): 234 | for infotype in self.capinfo: 235 | rows.append([infotype, self.capinfo[infotype]]) 236 | # Get the amount of bytes every 10 mins 237 | if self.get_bytes_histo(): 238 | rows.append(['Time Range (secs)', 'Amount of Packets' ]) 239 | for histo_header in sorted(self.histoinfo): 240 | rows.append([self.histoinfo[histo_header].split('|')[0], self.histoinfo[histo_header].split('|')[1]]) 241 | 242 | # Get more info from binetflow files 243 | elif 'binetflow' in self.get_type(): 244 | if self.get_binetflowinfos(): 245 | for infotype in self.binetflowinfo: 246 | rows.append([infotype, self.binetflowinfo[infotype]]) 247 | elif 'exe' in self.get_type(): 248 | # Get exe data 249 | # md5 250 | rows.append(['MD5', self.get_md5()]) 251 | 252 | print(table(header=['Key', 'Value'], rows=rows)) 253 | 254 | def __repr__(self): 255 | return('File id: {}. Name: {}. Type: {}'.format(self.get_id(), self.get_name(), self.get_type())) 256 | -------------------------------------------------------------------------------- /stf/core/models_constructors.py: -------------------------------------------------------------------------------- 1 | import persistent 2 | import BTrees.IOBTree 3 | from dateutil import parser 4 | import datetime 5 | 6 | from stf.common.out import * 7 | 8 | 9 | # Create one of these classes for each new model constructor you want to implement 10 | class Model_Constructor(object): 11 | """ 12 | The First Model constructor. Each of this objects is unique. We are going to instantiate them only once. 13 | Each model constructor object is related with a unique model. 14 | """ 15 | def __init__(self, id): 16 | self.id = id 17 | self.threshold_time_1 = False 18 | self.threshold_time_2 = False 19 | self.threshold_time_3 = False 20 | self.threshold_duration_1 = False 21 | self.threshold_duration_2 = False 22 | self.threshold_size_1 = False 23 | self.threshold_size_2 = False 24 | self.threshold_timeout = False 25 | self.use_multiples_timeouts = True 26 | self.models = {} 27 | 28 | def clean_models(self): 29 | self.models = {} 30 | 31 | def del_model(self, model_id): 32 | """ Delete this model from the list of models used by this constructor. This allow us to regenerate the state of a model without problems """ 33 | try: 34 | self.models.pop(model_id) 35 | except KeyError: 36 | print_error('There is no such model {} in the constructor to delete.'.format(model_id)) 37 | 38 | def set_name(self,name): 39 | self.name = name 40 | 41 | def set_description(self,description): 42 | self.description = description 43 | 44 | def get_state(self, flow, model_id): 45 | """ Receive the flow info and the model id and get a state""" 46 | # Temporal TD 47 | TD = -1 48 | state = '' 49 | 50 | # Get what we have about this model 51 | newtime = parser.parse(flow.get_starttime()) 52 | newsize = flow.get_totbytes() 53 | newduration = flow.get_duration() 54 | 55 | # This flow belongs to a known model, or is the first one? 56 | try: 57 | model = self.models[model_id] 58 | # We already have this model 59 | # Update T1. Just move T2 in there 60 | model['T1'] = model['T2'] 61 | # Get the new time from the new flow 62 | # Compute the new T2 63 | model['T2'] = newtime - model['LastTime'] 64 | # If T2 is negative, then we have an issue with the order of the file. Send an error and stop. The user should fix this, not us. 65 | if model['T2'].total_seconds() < 0: 66 | print_error('Model: {}'.format(model)) 67 | print_error('T1 is: {}'.format(model['T1'].total_seconds())) 68 | print_error('T2 is: {}'.format(model['T2'].total_seconds())) 69 | print_error('Flow new time is: {}'.format(newtime)) 70 | print_error('Flow last time is: {}'.format(model['LastTime'])) 71 | print_error('The last flow is: {}'.format(flow)) 72 | print_error('The binetflow file is not sorted. Please delete this file from the dataset, sort it (cat file.biargus |sort -n > newfile.biargus) and add it back. We can not modify a file on disk.') 73 | print_error('Flow: '.format(flow)) 74 | return False 75 | # Update the lasttime for next time 76 | model['LastTime'] = newtime 77 | except KeyError: 78 | # First time we see this model. Initialize the values 79 | self.models[model_id]={} 80 | self.models[model_id]['T1'] = False 81 | self.models[model_id]['T2'] = False 82 | self.models[model_id]['LastTime'] = newtime 83 | model = self.models[model_id] 84 | 85 | 86 | # We should get inside the next if only when T2 and T1 are not False. However, since also datatime(0) matches a False, we can only check to see if it is bool or not. 87 | # We are only using False when we start, so it is not necessary to check if it is False also. 88 | # Compute the periodicity 89 | if (isinstance(model['T1'], bool) and model['T1'] == False) or (isinstance(model['T2'], bool) and model['T2'] == False): 90 | periodic = -1 91 | elif not isinstance(model['T1'], bool) and not isinstance(model['T2'], bool): 92 | # We have some values. See which is larger 93 | try: 94 | if model['T2'] >= model['T1']: 95 | TD = datetime.timedelta(seconds=(model['T2'].total_seconds() / model['T1'].total_seconds())).total_seconds() 96 | else: 97 | TD = datetime.timedelta(seconds=(model['T1'].total_seconds() / model['T2'].total_seconds())).total_seconds() 98 | except ZeroDivisionError: 99 | #print_error('The time difference between flows was 0. Strange. We keep going anyway.') 100 | TD = 1 101 | # Decide the periodic based on TD and the thresholds 102 | if TD <= self.get_tt1(): 103 | # Strongly periodic 104 | periodic = 1 105 | elif TD < self.get_tt2(): 106 | # Weakly periodic 107 | periodic = 2 108 | elif TD < self.get_tt3(): 109 | # Weakly not periodic 110 | periodic = 3 111 | else: 112 | periodic = 4 113 | 114 | # Compute the duration 115 | if newduration <= self.get_td1(): 116 | duration = 1 117 | elif newduration > self.get_td1() and newduration <= self.get_td2(): 118 | duration = 2 119 | elif newduration > self.get_td2(): 120 | duration = 3 121 | 122 | # Compute the size 123 | if newsize <= self.get_ts1(): 124 | size = 1 125 | elif newsize > self.get_ts1() and newsize <= self.get_ts2(): 126 | size = 2 127 | elif newsize > self.get_ts2(): 128 | size = 3 129 | 130 | # Compute the state 131 | if periodic == -1: 132 | if size == 1: 133 | if duration == 1: 134 | state += '1' 135 | elif duration == 2: 136 | state += '2' 137 | elif duration == 3: 138 | state += '3' 139 | elif size == 2: 140 | if duration == 1: 141 | state += '4' 142 | elif duration == 2: 143 | state += '5' 144 | elif duration == 3: 145 | state += '6' 146 | elif size == 3: 147 | if duration == 1: 148 | state += '7' 149 | elif duration == 2: 150 | state += '8' 151 | elif duration == 3: 152 | state += '9' 153 | elif periodic == 1: 154 | if size == 1: 155 | if duration == 1: 156 | state += 'a' 157 | elif duration == 2: 158 | state += 'b' 159 | elif duration == 3: 160 | state += 'c' 161 | elif size == 2: 162 | if duration == 1: 163 | state += 'd' 164 | elif duration == 2: 165 | state += 'e' 166 | elif duration == 3: 167 | state += 'f' 168 | elif size == 3: 169 | if duration == 1: 170 | state += 'g' 171 | elif duration == 2: 172 | state += 'h' 173 | elif duration == 3: 174 | state += 'i' 175 | elif periodic == 2: 176 | if size == 1: 177 | if duration == 1: 178 | state += 'A' 179 | elif duration == 2: 180 | state += 'B' 181 | elif duration == 3: 182 | state += 'C' 183 | elif size == 2: 184 | if duration == 1: 185 | state += 'D' 186 | elif duration == 2: 187 | state += 'E' 188 | elif duration == 3: 189 | state += 'F' 190 | elif size == 3: 191 | if duration == 1: 192 | state += 'G' 193 | elif duration == 2: 194 | state += 'H' 195 | elif duration == 3: 196 | state += 'I' 197 | elif periodic == 3: 198 | if size == 1: 199 | if duration == 1: 200 | state += 'r' 201 | elif duration == 2: 202 | state += 's' 203 | elif duration == 3: 204 | state += 't' 205 | elif size == 2: 206 | if duration == 1: 207 | state += 'u' 208 | elif duration == 2: 209 | state += 'v' 210 | elif duration == 3: 211 | state += 'w' 212 | elif size == 3: 213 | if duration == 1: 214 | state += 'x' 215 | elif duration == 2: 216 | state += 'y' 217 | elif duration == 3: 218 | state += 'z' 219 | elif periodic == 4: 220 | if size == 1: 221 | if duration == 1: 222 | state += 'R' 223 | elif duration == 2: 224 | state += 'S' 225 | elif duration == 3: 226 | state += 'T' 227 | elif size == 2: 228 | if duration == 1: 229 | state += 'U' 230 | elif duration == 2: 231 | state += 'V' 232 | elif duration == 3: 233 | state += 'W' 234 | elif size == 3: 235 | if duration == 1: 236 | state += 'X' 237 | elif duration == 2: 238 | state += 'Y' 239 | elif duration == 3: 240 | state += 'Z' 241 | #print_info('Model: {}, T1: {}, T2: {}, TD:{}, Periodicity: {}, State: {}'.format(model_id, model['T1'], model['T2'], [TD.total_seconds() if not isinstance(TD,int) else -1], periodic, state)) 242 | 243 | # Compute the new letters for the time of the periodicity. 244 | if not isinstance(model['T2'], bool): 245 | if model['T2'] <= datetime.timedelta(seconds=5): 246 | state += '.' 247 | elif model['T2'] <= datetime.timedelta(seconds=60): 248 | state += ',' 249 | elif model['T2'] <= datetime.timedelta(seconds=300): 250 | state += '+' 251 | elif model['T2'] <= datetime.timedelta(seconds=3600): 252 | state += '*' 253 | elif model['T2'] >= self.get_tto(): 254 | # We convert it to int because we count the amount of complete hours that timeouted. The remaining time is not a timeout... 255 | t2_in_hours = model['T2'].total_seconds() / self.get_tto().total_seconds() 256 | # Should be int always 257 | for i in range(int(t2_in_hours)): 258 | state += '0' 259 | 260 | # We store permanently the T1, T2 and TD values on each flow, so we can later analyze it 261 | flow.set_t1(model['T1']) 262 | flow.set_t2(model['T2']) 263 | flow.set_td(TD) 264 | flow.set_state(state) 265 | 266 | return state 267 | 268 | 269 | def get_id(self): 270 | return self.id 271 | 272 | def get_name(self): 273 | return self.name 274 | 275 | def get_description(self): 276 | return self.description 277 | 278 | def get_tt1(self): 279 | return self.threshold_time_1 280 | 281 | def get_tt2(self): 282 | return self.threshold_time_2 283 | 284 | def get_tt3(self): 285 | return self.threshold_time_3 286 | 287 | def get_td1(self): 288 | return self.threshold_duration_1 289 | 290 | def get_td2(self): 291 | return self.threshold_duration_2 292 | 293 | def get_ts1(self): 294 | return self.threshold_size_1 295 | 296 | def get_ts2(self): 297 | return self.threshold_size_2 298 | 299 | def get_tto(self): 300 | return self.threshold_timeout 301 | 302 | def set_tt1(self, value): 303 | self.threshold_time_1 = value 304 | 305 | def set_tt2(self, value): 306 | self.threshold_time_2 = value 307 | 308 | def set_tt3(self, value): 309 | self.threshold_time_3 = value 310 | 311 | def set_td1(self, value): 312 | self.threshold_duration_1 = value 313 | 314 | def set_td2(self, value): 315 | self.threshold_duration_2 = value 316 | 317 | def set_ts1(self, value): 318 | self.threshold_size_1 = value 319 | 320 | def set_ts2(self, value): 321 | self.threshold_size_2 = value 322 | 323 | def set_tto(self, value): 324 | self.threshold_timeout = value 325 | 326 | def set_use_mutiples_timeouts(self, value): 327 | self.use_multiples_timeouts = value 328 | 329 | def get_use_mutiples_timeouts(self): 330 | try: 331 | return self.use_multiples_timeouts 332 | except AttributeError: 333 | # If there is no info, by default use True 334 | return True 335 | 336 | 337 | 338 | ######################### 339 | ######################### 340 | ######################### 341 | class Models_Constructors(persistent.Persistent): 342 | def __init__(self): 343 | """ This class holds all the different constructors of behavioral models based on states""" 344 | self.default_model_constructor = 1 345 | self.models_constructors = BTrees.IOBTree.BTree() 346 | 347 | # Reapeat this for each new constructor 348 | 349 | # Add the first model constructor 350 | first_model_constructor = Model_Constructor(0) 351 | first_model_constructor.set_tt1(float(1.05)) 352 | first_model_constructor.set_tt2(float(1.1)) 353 | first_model_constructor.set_tt3(float(5)) 354 | first_model_constructor.set_td1(float(0.1)) 355 | first_model_constructor.set_td2(float(10)) 356 | first_model_constructor.set_ts1(float(125)) 357 | first_model_constructor.set_ts2(float(1100)) 358 | first_model_constructor.set_tto(datetime.timedelta(seconds=3600)) 359 | first_model_constructor.use_multiples_timeouts = True 360 | first_model_constructor.set_name('Model 0') 361 | first_model_constructor.set_description('To try at the thresholds.') 362 | self.models_constructors[first_model_constructor.get_id()] = first_model_constructor 363 | 364 | # Add the second model constructor 365 | second_model_constructor = Model_Constructor(1) 366 | second_model_constructor.set_tt1(float(1.05)) 367 | second_model_constructor.set_tt2(float(1.3)) 368 | second_model_constructor.set_tt3(float(5)) 369 | second_model_constructor.set_td1(float(0.1)) 370 | second_model_constructor.set_td2(float(10)) 371 | second_model_constructor.set_ts1(float(250)) 372 | second_model_constructor.set_ts2(float(1100)) 373 | second_model_constructor.set_tto(datetime.timedelta(seconds=3600)) 374 | second_model_constructor.use_multiples_timeouts = True 375 | second_model_constructor.set_name('Model Bundchen') 376 | second_model_constructor.set_description('Uses the symbols between flows to store the time. Better thresholds.') 377 | self.models_constructors[second_model_constructor.get_id()] = second_model_constructor 378 | 379 | # Add the third model constructor 380 | third_model_constructor = Model_Constructor(2) 381 | third_model_constructor.set_tt1(float(1.05)) 382 | third_model_constructor.set_tt2(float(1.3)) 383 | third_model_constructor.set_tt3(float(5)) 384 | third_model_constructor.set_td1(float(0.1)) 385 | third_model_constructor.set_td2(float(10)) 386 | third_model_constructor.set_ts1(float(250)) 387 | third_model_constructor.set_ts2(float(1100)) 388 | third_model_constructor.set_tto(datetime.timedelta(seconds=3600)) 389 | third_model_constructor.use_multiples_timeouts = False 390 | third_model_constructor.set_name('Model Moss') 391 | third_model_constructor.set_description('Uses the symbols between flows to store the time. Better thresholds.') 392 | self.models_constructors[third_model_constructor.get_id()] = third_model_constructor 393 | 394 | def has_constructor_id(self, constructor_id): 395 | try: 396 | t = self.models_constructors[constructor_id] 397 | return True 398 | except KeyError: 399 | return False 400 | 401 | def get_constructor(self, id): 402 | """ Return the constructors ids""" 403 | return self.models_constructors[id] 404 | 405 | def get_constructors_ids(self): 406 | """ Return the constructors ids""" 407 | return self.models_constructors.keys() 408 | 409 | def get_default_constructor(self): 410 | """ Since we return an object, all the models share the same constructor """ 411 | return self.models_constructors[self.default_model_constructor] 412 | 413 | def list_constructors(self): 414 | print_info('List of all the models constructors available') 415 | rows = [] 416 | for constructor in self.models_constructors.values(): 417 | rows.append([constructor.get_name(), constructor.get_id(), constructor.get_description(), constructor.get_tt1(), constructor.get_tt2(), constructor.get_tt3(), constructor.get_td1(), constructor.get_td2(), constructor.get_ts1(), constructor.get_ts2(), constructor.get_tto()]) 418 | print(table(header=['Name', 'Id', 'Description', 'Tt1', 'Tt2', 'Tt3', 'Td1', 'Td2', 'Ts1', 'Ts2', 'Tto'], rows=rows)) 419 | 420 | 421 | 422 | __modelsconstructors__ = Models_Constructors() 423 | -------------------------------------------------------------------------------- /stf/core/notes.py: -------------------------------------------------------------------------------- 1 | import persistent 2 | import BTrees.IOBTree 3 | import tempfile 4 | import os 5 | from datetime import datetime 6 | from subprocess import Popen, PIPE 7 | 8 | from stf.common.out import * 9 | 10 | ############################### 11 | ############################### 12 | ############################### 13 | class Note(persistent.Persistent): 14 | """ 15 | The Note 16 | """ 17 | def __init__(self, id): 18 | self.id = id 19 | self.text = "" 20 | 21 | def get_id(self): 22 | return self.id 23 | 24 | def edit(self): 25 | # Create a new temporary file. 26 | tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.md') 27 | # Write the text that we have into the temp file 28 | tmp.write(self.text) 29 | tmp.file.flush() 30 | # Open the temporary file with the default editor, or with nano. 31 | os.system('"${EDITOR:-nano}" ' + tmp.name) 32 | # Go to the beginning of the file 33 | tmp.file.flush() 34 | tmp.file.seek(0) 35 | self.text = tmp.file.read() 36 | # Finally, remove the temporary file. 37 | os.remove(tmp.name) 38 | 39 | def get_text(self): 40 | return self.text 41 | 42 | def delete_text(self): 43 | """ Delete the text of the note """ 44 | self.text = "" 45 | 46 | def get_note(self): 47 | """ Return text of the note """ 48 | return self.text 49 | 50 | def get_short_note(self): 51 | """ Return text of the note until the first enter""" 52 | try: 53 | enter = self.text.index('\n') 54 | return self.text[:enter] 55 | except ValueError: 56 | # There is no enter 57 | return self.text[0:80] 58 | 59 | def __repr__(self): 60 | return self.text 61 | 62 | def add_text(self, text_to_add): 63 | """ Add text to the note without intervention """ 64 | self.text += text_to_add 65 | 66 | def show_text(self): 67 | f = tempfile.NamedTemporaryFile() 68 | f.write(self.text) 69 | f.flush() 70 | p = Popen('less -R ' + f.name, shell=True, stdin=PIPE) 71 | p.communicate() 72 | sys.stdout = sys.__stdout__ 73 | f.close() 74 | 75 | def has_text(self, text_to_search): 76 | """ Searchs a text. Ignore case""" 77 | if text_to_search.lower() in self.text.lower(): 78 | return True 79 | else: 80 | return False 81 | 82 | 83 | ############################### 84 | ############################### 85 | ############################### 86 | class Group_of_Notes(persistent.Persistent): 87 | """ This class holds all the notes""" 88 | def __init__(self): 89 | self.notes = BTrees.IOBTree.BTree() 90 | # We are not storing here the relationship between the note and the object to which the note is related. The relationship is stored in the other object. 91 | 92 | def get_note(self, note_id): 93 | """ Return all the notes """ 94 | try: 95 | return self.notes[note_id] 96 | except KeyError: 97 | return False 98 | 99 | def get_notes_ids(self): 100 | """ Return all the notes ids""" 101 | return self.notes.keys() 102 | 103 | def get_notes(self): 104 | """ Return all the notes """ 105 | return self.notes.values() 106 | 107 | def new_note(self): 108 | """ Creates a new note and returns its id """ 109 | # Get the new id for this note 110 | try: 111 | # Get the id of the last note in the database 112 | note_id = self.notes[list(self.notes.keys())[-1]].get_id() + 1 113 | except (KeyError, IndexError): 114 | note_id = 0 115 | new_note = Note(note_id) 116 | # Store it 117 | self.notes[note_id] = new_note 118 | return note_id 119 | 120 | def delete_note(self, note_id): 121 | try: 122 | # Just in case delete the text of the note before 123 | note = self.notes[note_id] 124 | note.delete_text() 125 | # Delete the note 126 | self.notes.pop(note_id) 127 | except KeyError: 128 | print_error('No such note id.') 129 | 130 | def get_short_note(self, note_id): 131 | try: 132 | note = self.notes[note_id] 133 | return note.get_short_note() 134 | except KeyError: 135 | return '' 136 | 137 | def list_notes(self): 138 | """ List all the notes """ 139 | f = tempfile.NamedTemporaryFile() 140 | for note in self.get_notes(): 141 | f.write(cyan('Note {}'.format(note.get_id())) + '\n') 142 | f.write(note.get_text() + '\n') 143 | f.flush() 144 | p = Popen('less -R ' + f.name, shell=True, stdin=PIPE) 145 | p.communicate() 146 | sys.stdout = sys.__stdout__ 147 | f.close() 148 | 149 | def add_auto_text_to_note(self, note_id, text_to_add): 150 | """ Gets a text to be automatically added to the note. Used to log internal operations of the framework in the notes. Such as, the flows in this connection had been trimed """ 151 | note = self.get_note(note_id) 152 | if note: 153 | now = str(datetime.now()) 154 | note.add_text('\n[#] ' + now + ': ' + text_to_add) 155 | 156 | def show_note(self, note_id): 157 | """ Show a note """ 158 | note = self.get_note(note_id) 159 | if note: 160 | note.show_text() 161 | 162 | def edit_note(self, note_id): 163 | """ Edit a note """ 164 | note = self.get_note(note_id) 165 | if note: 166 | note.edit() 167 | else: 168 | print_error('No such note id.') 169 | 170 | def search_text(self, text_to_search): 171 | """ Search a text in all notes """ 172 | for note in self.get_notes(): 173 | if note.has_text(text_to_search): 174 | print_info(cyan('Note {}'.format(note.get_id()))) 175 | print'{}'.format(note.get_text()) 176 | 177 | __notes__ = Group_of_Notes() 178 | -------------------------------------------------------------------------------- /stf/core/plugins.py: -------------------------------------------------------------------------------- 1 | # This file is part of Viper - https://github.com/botherder/viper 2 | # See the file 'LICENSE' for copying permission. 3 | 4 | import pkgutil 5 | import inspect 6 | 7 | from stf.common.out import print_warning 8 | from stf.common.abstracts import Module 9 | 10 | def load_modules(): 11 | # Import modules package. 12 | # This loads the modules in the modules folder. Is very relative to the starting position of the stf 13 | import modules 14 | 15 | plugins = dict() 16 | 17 | # Walk recursively through all modules and packages. 18 | for loader, module_name, ispkg in pkgutil.walk_packages(modules.__path__, modules.__name__ + '.'): 19 | # If current item is a package, skip. 20 | if ispkg: 21 | continue 22 | 23 | # Try to import the module, otherwise skip. 24 | try: 25 | module = __import__(module_name, globals(), locals(), ['dummy'], -1) 26 | except ImportError as e: 27 | print_warning("Something wrong happened while importing the module {0}: {1}".format(module_name, e)) 28 | continue 29 | 30 | # Walk through all members of currently imported modules. 31 | for member_name, member_object in inspect.getmembers(module): 32 | # Check if current member is a class. 33 | if inspect.isclass(member_object): 34 | # Yield the class if it's a subclass of Module. 35 | if issubclass(member_object, Module) and member_object is not Module: 36 | plugins[member_object.cmd] = dict(obj=member_object, description=member_object.description) 37 | 38 | return plugins 39 | 40 | __modules__ = load_modules() 41 | -------------------------------------------------------------------------------- /stf/core/ui/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/stratosphereips/StratosphereTestingFramework/381eca74358dfa3847727314f2ae615cd5c8170a/stf/core/ui/__init__.py -------------------------------------------------------------------------------- /stf/core/ui/commands.py: -------------------------------------------------------------------------------- 1 | # This file is part of the Stratosphere Testing Framework 2 | # See the file 'LICENSE' for copying permission. 3 | # Most of this file is copied from Viper 4 | 5 | import argparse 6 | import os 7 | import time 8 | import fnmatch 9 | import tempfile 10 | import shutil 11 | import transaction 12 | 13 | from stf.common.out import * 14 | from stf.core.dataset import __datasets__ 15 | #from stf.core.experiment import __experiments__ 16 | from stf.core.database import __database__ 17 | from stf.core.connections import __group_of_group_of_connections__ 18 | from stf.core.models_constructors import __modelsconstructors__ 19 | from stf.core.models import __groupofgroupofmodels__ 20 | from stf.core.notes import __notes__ 21 | from stf.core.labels import __group_of_labels__ 22 | from stf.core.plugins import __modules__ 23 | 24 | class Commands(object): 25 | 26 | def __init__(self): 27 | # Map commands to their related functions. 28 | self.commands = dict( 29 | help=dict(obj=self.cmd_help, description="Show this help message"), 30 | #info=dict(obj=self.cmd_info, description="Show information on the opened experiment"), 31 | clear=dict(obj=self.cmd_clear, description="Clear the console"), 32 | #experiments=dict(obj=self.cmd_experiments, description="List or switch to existing experiments"), 33 | datasets=dict(obj=self.cmd_datasets, description="Manage the datasets"), 34 | connections=dict(obj=self.cmd_connections, description="Manage the connections. A dataset should be selected first."), 35 | models=dict(obj=self.cmd_models, description="Manage the models. A dataset should be selected first."), 36 | database=dict(obj=self.cmd_database, description="Manage the database."), 37 | notes=dict(obj=self.cmd_notes, description="Manage the notes."), 38 | labels=dict(obj=self.cmd_labels, description="Manage the labels."), 39 | exit=dict(obj=self.cmd_exit, description="Exit"), 40 | ) 41 | 42 | ## 43 | # CLEAR 44 | # 45 | # This command simply clears the shell. 46 | def cmd_clear(self, *args): 47 | os.system('clear') 48 | 49 | ## 50 | # HELP 51 | # 52 | # This command simply prints the help message. 53 | # It lists both embedded commands and loaded modules. 54 | def cmd_help(self, *args): 55 | print(bold("Commands:")) 56 | 57 | rows = [] 58 | for command_name, command_item in self.commands.items(): 59 | rows.append([command_name, command_item['description']]) 60 | 61 | #rows.append(["exit, quit", "Exit Viper"]) 62 | rows = sorted(rows, key=lambda entry: entry[0]) 63 | 64 | print(table(['Command', 'Description'], rows)) 65 | print("") 66 | print(bold("Modules:")) 67 | 68 | rows = [] 69 | for module_name, module_item in __modules__.items(): 70 | rows.append([module_name, module_item['description']]) 71 | 72 | rows = sorted(rows, key=lambda entry: entry[0]) 73 | 74 | print(table(['Command', 'Description'], rows)) 75 | 76 | ## 77 | # NOTES 78 | # 79 | # This command works with notes 80 | def cmd_notes(self, *args): 81 | parser = argparse.ArgumentParser(prog="notes", description="Manage notes", epilog="Manage notes") 82 | parser.add_argument('-l', '--listnotes', action="store_true", help="List all the notes in the system.") 83 | parser.add_argument('-d', '--deletenote', metavar="note_id", help="Delete a note.") 84 | parser.add_argument('-f', '--filter', metavar="filter", nargs = '+', help="Use this filter to work with notes. You can use multiple filter separated by a space. Format: \"variable[=<>]value\". You can use the variables: text. For example: -f text=hi text!=p2p.") 85 | parser.add_argument('-s', '--show', type=int, metavar="note_id", help="Show this note id.") 86 | parser.add_argument('-e', '--edit', type=int, metavar="note_id", help="Edit this note id.") 87 | parser.add_argument('-S', '--search', type=str, metavar="text", help="Search a text in all the notes, and list the notes.") 88 | 89 | 90 | try: 91 | args = parser.parse_args(args) 92 | except: 93 | return 94 | 95 | # Subcomand to list the notes 96 | if args.listnotes: 97 | __notes__.list_notes() 98 | 99 | # Subcomand to delte a note 100 | elif args.deletenote: 101 | __notes__.delete_note(int(args.deletenote)) 102 | __database__.root._p_changed = True 103 | 104 | # Subcomand to show a note 105 | elif args.show: 106 | __notes__.show_note(args.show) 107 | 108 | # Subcomand to edit a note 109 | elif args.edit >= 0: 110 | __notes__.edit_note(args.edit) 111 | __database__.root._p_changed = True 112 | 113 | # Subcomand to search a text 114 | elif args.search: 115 | __notes__.search_text(args.search) 116 | 117 | ## 118 | # MODELS 119 | # 120 | # This command works with models 121 | def cmd_models(self, *args): 122 | parser = argparse.ArgumentParser(prog="models", description="Manage models", epilog="Manage models") 123 | parser.add_argument('-s', '--listconstructors', action="store_true", help="List all models constructors available.") 124 | parser.add_argument('-l', '--listgroups', action="store_true", help="List all the groups of models. If a dataset is selected it only shows the models in that dataset.") 125 | parser.add_argument('-g', '--generate', action="store_true", help="Generate the models for the current dataset.") 126 | parser.add_argument('-d', '--deletegroup', metavar="group_model_id", help="Delete a group of models.") 127 | parser.add_argument('-D', '--deletemodel', metavar="group_model_id", help="Delete a model (4tuple) from this group. With -D give the id of the group. Use -i to give the model id to delete (4-tuple) or -f to use a filter.") 128 | parser.add_argument('-i', '--modelid', metavar="model_id", help="Use this model id (4-tuple). Commonly used with -D.") 129 | parser.add_argument('-L', '--listmodels', metavar="group_model_id", help="List the models inside a group. You can use filters.") 130 | parser.add_argument('-C', '--countmodels', metavar="group_model_id", help="Count the models inside a group.") 131 | parser.add_argument('-f', '--filter', metavar="filter", nargs = '+', default="", help="Use this filter to work with models. You can use multiple filter separated by a space. Format: \"variable[=<>]value\". You can use the variables: statelength, name and labelname. For example: -f statelength>100 name=tcp. Another example: -f name=-tcp- labelname=Botnet") 132 | parser.add_argument('-H', '--histogram', metavar="group_model_id", help="Plot a histogram of the lengths of models states in the given id of group of models.") 133 | parser.add_argument('-N', '--delnote', metavar='group_model_id', help="Delete completely the note related with this model id. Use -i to give the model id to add the note to (4-tuple).") 134 | parser.add_argument('-n', '--editnote', metavar='group_model_id', help="Edit the note related with this model id. Use -i to give the model id to add the note to (4-tuple).") 135 | parser.add_argument('-o', '--listnotes', default=0, metavar='group_model_id', help="List the notes related with this model id. You can use the -f with filters here.") 136 | parser.add_argument('-a', '--amountoflettersinstate', default=0, metavar='amount_of_letters', help="When used with -L, limit the maximum amount of letters in the state to show per line. Helps avoiding dangerously long lines.") 137 | parser.add_argument('-c', '--constructor', metavar="constructor_id", type=int, help="Use this constructor for generating the new models. Use optionally with -g.") 138 | parser.add_argument('-e', '--exportasciimodels', metavar="group_model_id", help="Export an ascii list of the all the connections, labels and letters in this Model id. Useful for external analysis.") 139 | 140 | 141 | try: 142 | args = parser.parse_args(args) 143 | except: 144 | return 145 | 146 | # Subcomand to list the constructors 147 | if args.listconstructors: 148 | __modelsconstructors__.list_constructors() 149 | 150 | # Subcomand to list the models 151 | elif args.listgroups: 152 | __groupofgroupofmodels__.list_groups() 153 | 154 | # Subcomand to generate the models 155 | elif args.generate: 156 | if args.constructor != None: 157 | if __modelsconstructors__.has_constructor_id(args.constructor): 158 | constructor = int(args.constructor) 159 | else: 160 | print_error('No such constructor id available.') 161 | return False 162 | else: 163 | constructor = __modelsconstructors__.get_default_constructor().get_id() 164 | __groupofgroupofmodels__.generate_group_of_models(constructor) 165 | __database__.root._p_changed = True 166 | 167 | # Subcomand to delete the group of models of the current dataset 168 | elif args.deletegroup: 169 | __groupofgroupofmodels__.delete_group_of_models(args.deletegroup) 170 | __database__.root._p_changed = True 171 | 172 | # Subcomand to list the models in a group 173 | elif args.listmodels: 174 | __groupofgroupofmodels__.list_models_in_group(args.listmodels, args.filter, int(args.amountoflettersinstate)) 175 | 176 | # Subcommand to export the ascii of the models 177 | elif args.exportasciimodels: 178 | __groupofgroupofmodels__.export_models_in_group(args.exportasciimodels, args.filter) 179 | 180 | # Subcomand to delete a model from a group by id or filter 181 | elif args.deletemodel: 182 | # By id or filter? 183 | if args.modelid: 184 | # By id 185 | __groupofgroupofmodels__.delete_a_model_from_the_group_by_id(args.deletemodel, args.modelid) 186 | __database__.root._p_changed = True 187 | elif args.filter: 188 | # By filter 189 | __groupofgroupofmodels__.delete_a_model_from_the_group_by_filter(args.deletemodel, args.filter) 190 | __database__.root._p_changed = True 191 | else: 192 | print_error('You should provide the id of the model (4-tuple) with -i or a filter with -f') 193 | 194 | # Subcomand to count the amount of models 195 | elif args.countmodels: 196 | __groupofgroupofmodels__.count_models_in_group(args.countmodels, args.filter) 197 | __database__.root._p_changed = True 198 | 199 | # Subcomand to plot histogram of states lengths 200 | elif args.histogram: 201 | __groupofgroupofmodels__.plot_histogram(args.histogram, args.filter) 202 | 203 | # Subcomand to edit the note of this model 204 | elif args.editnote: 205 | if args.modelid: 206 | __groupofgroupofmodels__.edit_note(args.editnote, args.modelid) 207 | __database__.root._p_changed = True 208 | else: 209 | print_error('You should give a model id also with -i.') 210 | 211 | # Subcomand to delete the note of this model 212 | elif args.delnote : 213 | if args.modelid: 214 | __groupofgroupofmodels__.del_note(args.delnote, args.modelid) 215 | __database__.root._p_changed = True 216 | else: 217 | print_error('You should give a model id also with -i.') 218 | 219 | # Subcomand to list the note of this model 220 | elif args.listnotes : 221 | __groupofgroupofmodels__.list_notes(args.listnotes, args.filter) 222 | __database__.root._p_changed = True 223 | 224 | ## 225 | # CONNECTIONS 226 | # 227 | # This command works with connections 228 | def cmd_connections(self, *args): 229 | parser = argparse.ArgumentParser(prog="connections", description="Manage connections", epilog="Manage connections") 230 | parser.add_argument('-l', '--list', action="store_true", help="List all existing connections") 231 | parser.add_argument('-g', '--generate', action="store_true", help="Generate the connections from the binetflow file in the current dataset") 232 | parser.add_argument('-d', '--delete', metavar="group_of_connections_id", help="Delete the group of connections.") 233 | parser.add_argument('-L', '--listconnections', metavar="group_connection_id", help="List the connections inside a group.") 234 | parser.add_argument('-F', '--showflows', metavar="connection_id", type=str, help="List the flows inside a specific connection.") 235 | parser.add_argument('-f', '--filter', metavar="filter", nargs='+', help="Use this filter to work with connections. Format: \"variable[!=<>]value\". You can use the variables: name, flowamount. Example: \"name=tcp\". Or \"flowamount<10\"") 236 | parser.add_argument('-D', '--deleteconnection', metavar="group_connection_id", help="Delete a connection from the group. This is the id of the group. Use -i to give the connection id to delete (4-tuple) or -f to use a filter.") 237 | parser.add_argument('-i', '--connectionid', metavar="connection_id", help="Use this connection id (4-tuple). Commonly used with -D.") 238 | parser.add_argument('-M', '--deleteconnectionifmodel', metavar="group_connection_id", help="Delete the connections from the group which models were deleted. Only give the connection group id. Useful to clean the database of connections that are not used.") 239 | parser.add_argument('-t', '--trimgroupid', metavar="group_connection_id", help="Trim all the connections so that each connection has at most 100 flows. Only give the connection group id. Useful to have some info about the connections but not all the data.") 240 | parser.add_argument('-a', '--amounttotrim', metavar="amount_to_trim", type=int, help="Define the amount of flows to trim with -t. By default 100.") 241 | parser.add_argument('-C', '--countconnections', metavar="group_connection_id", help="Count the amount of connections matching the filter. This is the id of the group.") 242 | parser.add_argument('-H', '--histogram', metavar="connection_id", type=str, help="Show the histograms for state len, duration and size of all the flows in this connection id (4-tuple).") 243 | try: 244 | args = parser.parse_args(args) 245 | except: 246 | return 247 | 248 | # Subcomand to list 249 | if args.list: 250 | __group_of_group_of_connections__.list_group_of_connections() 251 | 252 | # Subcomand to create a new group of connections 253 | elif args.generate: 254 | __group_of_group_of_connections__.create_group_of_connections() 255 | __database__.root._p_changed = True 256 | 257 | # Subcomand to delete a group of connections 258 | elif args.delete: 259 | __group_of_group_of_connections__.delete_group_of_connections(int(args.delete)) 260 | __database__.root._p_changed = True 261 | 262 | # Subcomand to list the connections in a group 263 | elif args.listconnections: 264 | filter = '' 265 | try: 266 | filter = args.filter 267 | except AttributeError: 268 | pass 269 | try: 270 | __group_of_group_of_connections__.list_connections_in_group(int(args.listconnections), filter) 271 | except ValueError: 272 | print_error('The id of the group of connections should be an integer.') 273 | __database__.root._p_changed = True 274 | 275 | # Subcomand to show the flows in a connection 276 | elif args.showflows: 277 | filter = '' 278 | try: 279 | filter = args.filter 280 | except AttributeError: 281 | pass 282 | __group_of_group_of_connections__.show_flows_in_connnection(args.showflows, filter) 283 | __database__.root._p_changed = True 284 | 285 | # Subcomand to delete a connection from a group by id 286 | elif args.deleteconnection: 287 | if args.connectionid: 288 | # By id 289 | __group_of_group_of_connections__.delete_a_connection_from_the_group_by_id(args.deleteconnection, args.connectionid) 290 | elif args.filter: 291 | __group_of_group_of_connections__.delete_a_connection_from_the_group_by_filter(args.deleteconnection, args.filter) 292 | __database__.root._p_changed = True 293 | 294 | # Subcomand to delete the connections from a group which models were deleted 295 | elif args.deleteconnectionifmodel: 296 | __group_of_group_of_connections__.delete_a_connection_from_the_group_if_model_deleted(int(args.deleteconnectionifmodel)) 297 | __database__.root._p_changed = True 298 | 299 | # Subcomand to trim the amount of flows in the connections 300 | elif args.trimgroupid: 301 | # Now just trim to keep 100 flows 302 | if args.amounttotrim: 303 | amount_to_trim = args.amounttotrim 304 | else: 305 | amount_to_trim = 100 306 | __group_of_group_of_connections__.trim_flows(int(args.trimgroupid), amount_to_trim) 307 | __database__.root._p_changed = True 308 | 309 | # Subcomand to count the amount of models 310 | elif args.countconnections: 311 | try: 312 | filter = args.filter 313 | except AttributeError: 314 | pass 315 | __group_of_group_of_connections__.count_connections_in_group(args.countconnections, filter) 316 | __database__.root._p_changed = True 317 | 318 | # Subcomand to show the histograms 319 | elif args.histogram: 320 | __group_of_group_of_connections__.show_histograms(args.histogram) 321 | 322 | ## 323 | # DATASETS 324 | # 325 | # This command works with datasets 326 | def cmd_datasets(self, *args): 327 | parser = argparse.ArgumentParser(prog="datasets", description="Manage datasets", epilog="Manage datasets") 328 | group = parser.add_mutually_exclusive_group() 329 | group.add_argument('-l', '--list', action="store_true", help="List all existing datasets.") 330 | group.add_argument('-c', '--create', metavar='filename', help="Create a new dataset from a file.") 331 | group.add_argument('-d', '--delete', metavar='dataset_id', help="Delete a dataset.") 332 | group.add_argument('-s', '--select', metavar='dataset_id', help="Select a dataset to work with. Enables the following commands on the dataset.") 333 | group.add_argument('-f', '--list_files', action='store_true', help="List all the files in the current dataset") 334 | group.add_argument('-F', '--file', metavar='file_id', help="Give more info about the selected file in the current dataset.") 335 | group.add_argument('-a', '--add', metavar='file_id', help="Add a file to the current dataset.") 336 | group.add_argument('-D', '--dele', metavar='file_id', help="Delete a file from the dataset.") 337 | group.add_argument('-g', '--generate', action='store_true', help="Try to generate the biargus and binetflow files for the selected dataset if they do not exists.") 338 | group.add_argument('-u', '--unselect', action='store_true', help="Unselect the current dataset.") 339 | group.add_argument('-n', '--editnote', metavar='dataset_id', help="Edit the note related with this dataset id.") 340 | group.add_argument('-N', '--delnote', metavar='dataset_id', help="Delete completely the note related with this dataset id.") 341 | group.add_argument('-o', '--editfolder', metavar='dataset_id', type=str, help="Edit the main folder of this dataset. Useful when you upload files from different machines and then you move them around.") 342 | 343 | try: 344 | args = parser.parse_args(args) 345 | except: 346 | return 347 | 348 | # Subcomand to list 349 | if args.list: 350 | __datasets__.list() 351 | 352 | # Subcomand to create 353 | elif args.create: 354 | __datasets__.create(args.create) 355 | __database__.root._p_changed = True 356 | 357 | # Subcomand to delete 358 | elif args.delete: 359 | __datasets__.delete(int(args.delete)) 360 | __database__.root._p_changed = True 361 | 362 | # Subcomand to select a dataset 363 | elif args.select : 364 | __datasets__.select(args.select) 365 | 366 | # Subcomand to list files 367 | elif args.list_files: 368 | __datasets__.list_files() 369 | __database__.root._p_changed = True 370 | 371 | # Subcomand to get info about a file in a dataset 372 | elif args.file : 373 | __datasets__.info_about_file(args.file) 374 | __database__.root._p_changed = True 375 | 376 | # Subcomand to add a file to the dataset 377 | elif args.add : 378 | __datasets__.add_file(args.add) 379 | __database__.root._p_changed = True 380 | 381 | # Subcomand to delete a file from the dataset 382 | elif args.dele : 383 | __datasets__.del_file(args.dele) 384 | __database__.root._p_changed = True 385 | 386 | # Subcomand to generate the biargus and binetflow files in a dataset 387 | elif args.generate : 388 | __datasets__.generate_argus_files() 389 | __database__.root._p_changed = True 390 | 391 | # Subcomand to unselect the current dataset 392 | elif args.unselect : 393 | __datasets__.unselect_current() 394 | 395 | # Subcomand to edit the note of this dataset 396 | elif args.editnote : 397 | __datasets__.edit_note(args.editnote) 398 | __database__.root._p_changed = True 399 | 400 | # Subcomand to delete the note of this dataset 401 | elif args.delnote : 402 | __datasets__.del_note(args.delnote) 403 | __database__.root._p_changed = True 404 | 405 | # Subcomand to edit the folder of this dataset 406 | elif args.editfolder : 407 | __datasets__.edit_folder(args.editfolder) 408 | __database__.root._p_changed = True 409 | else: 410 | parser.print_usage() 411 | 412 | 413 | ## 414 | # DATABASE 415 | # 416 | def cmd_database(self, *args): 417 | parser = argparse.ArgumentParser(prog="database", description="Manage the database", epilog="Manage database") 418 | group = parser.add_mutually_exclusive_group() 419 | group.add_argument('-i', '--info', action="store_true", help="Info about the database connection") 420 | group.add_argument('-r', '--revert', action="store_true", help="Revert the connection of the database to the state before the last pack") 421 | group.add_argument('-p', '--pack', action="store_true", help="Pack the database") 422 | group.add_argument('-c', '--commit', action="store_true", help="Commit the changes") 423 | group.add_argument('-l', '--list', action="store_true", help="List the structures in the db.") 424 | group.add_argument('-d', '--delete', metavar="structurename", help="Delete the given structure from the db. Specify the complete name.") 425 | 426 | try: 427 | args = parser.parse_args(args) 428 | except: 429 | return 430 | 431 | # Subcomand to get info 432 | if args.info: 433 | __database__.info() 434 | 435 | # Subcomand to delete a structures 436 | elif args.delete: 437 | __database__.delete_structure(args.delete) 438 | 439 | # Subcomand to list the structures 440 | elif args.list: 441 | __database__.list_structures() 442 | 443 | # Subcomand to revert the database 444 | elif args.revert: 445 | __database__.revert() 446 | 447 | # Subcomand to pack he database 448 | elif args.pack: 449 | __database__.pack() 450 | 451 | # Subcomand to commit the changes 452 | elif args.commit: 453 | __database__.commit() 454 | 455 | 456 | ## 457 | # LABELS 458 | # 459 | # This command works with labels 460 | def cmd_labels(self, *args): 461 | parser = argparse.ArgumentParser(prog="labels", description="Manage labels", epilog="Manage labels") 462 | parser.add_argument('-l', '--list', action="store_true", help="List all existing labels.") 463 | parser.add_argument('-a', '--add', action="store_true", help="Add a label. Use -c to add to a connection_id (or IP) or -f to add to a group of connections id.") 464 | parser.add_argument('-c', '--connectionid', metavar="connection_id", help="Together with -a, add a label to the given connection_id or IP address. You should use -g to specify the id of the group of models.") 465 | parser.add_argument('-d', '--delete', metavar="label_id", help="Delete a label given the label number id. If you use a dash you can delete ranges of label ids. For example: -d 20-30") 466 | parser.add_argument('-F', '--deletefilter', action="store_true", help="Delete labels using the filter given with -f.") 467 | parser.add_argument('-D', '--deleteconnection', metavar="connection_id", help="Delete a label given a connection id to delete (4-tuple). You must give the group of model id with -g.") 468 | parser.add_argument('-g', '--modelgroupid', metavar="modelgroupid", help="Id of the group of models. Used with -a.") 469 | parser.add_argument('-m', '--migrate', action="store_true", help="Migrate <= 0.1.2alpha labels to the new database.") 470 | parser.add_argument('-f', '--filter', metavar="filter", nargs='+', default="", help="Use this filter to work with labels. Format: \"variable[!=<>]value\". You can use the variables: name, id, groupid and connid. Example: \"name=Botnet\". If you use -f to add labels, you should also specify -g. the variable connid is only used to assign a label to multiple connections") 471 | 472 | try: 473 | args = parser.parse_args(args) 474 | except: 475 | return 476 | 477 | # Subcomand to list labels 478 | if args.list: 479 | __group_of_labels__.list_labels(args.filter) 480 | 481 | # Subcomand to add a label 482 | elif args.add: 483 | if args.modelgroupid: 484 | # To a group of connections id using filters 485 | if args.filter: 486 | __group_of_labels__.add_label_with_filter(args.modelgroupid, args.filter) 487 | # To a unique connections id 488 | elif args.connectionid: 489 | __group_of_labels__.add_label(args.modelgroupid, args.connectionid) 490 | else: 491 | print_error('Please specify the id of the group of models where this connection belongs with -g.') 492 | 493 | # Subcomand to delete a label 494 | elif args.delete: 495 | __group_of_labels__.del_label(args.delete) 496 | # Subcomand to delete a label with filter 497 | elif args.deletefilter: 498 | if args.filter: 499 | __group_of_labels__.del_label(False, args.filter) 500 | # Subcomand to delete a specific connection 501 | elif args.deleteconnection: 502 | if args.modelgroupid: 503 | __group_of_labels__.delete_connection(args.modelgroupid, args.deleteconnection) 504 | else: 505 | print_error('You should give a group of models id with -g.') 506 | 507 | # Subcomand to migrate old labels 508 | elif args.migrate: 509 | __group_of_labels__.migrate_old_labels() 510 | 511 | 512 | ## 513 | # EXIT 514 | # 515 | def cmd_exit(self): 516 | # Exit is handled in other place. This is so it can appear in the autocompletion 517 | pass 518 | 519 | -------------------------------------------------------------------------------- /stf/core/ui/console.py: -------------------------------------------------------------------------------- 1 | # This file is part of the Stratosphere Testing Framework 2 | # See the file 'LICENSE' for copying permission. 3 | # A large part of this file is taken from the Viper tool 4 | 5 | import os 6 | import glob 7 | import atexit 8 | import readline 9 | import traceback 10 | 11 | from stf.common.out import * 12 | from stf.core.plugins import __modules__ 13 | 14 | 15 | version = "0.1.6alpha" 16 | 17 | def logo(): 18 | print(""" 19 | Stratosphere Testing Framework 20 | https://stratosphereips.org 21 | _ __ 22 | | | / _| 23 | ___| |_| |_ 24 | / __| __| _| 25 | \__ \ |_| | 26 | ... |___/\__|_| ... 27 | """+version+""" 28 | 29 | """) 30 | 31 | 32 | 33 | class Console(object): 34 | 35 | def __init__(self): 36 | # Create the nessesary folders first 37 | self.create_folders() 38 | 39 | # I have to move the import here. 40 | from stf.core.ui.commands import Commands 41 | 42 | # From some reason we should initialize the db from a method, we can not do it in the constructor 43 | from stf.core.database import __database__ 44 | __database__.start() 45 | 46 | # This will keep the main loop active as long as it's set to True. 47 | self.active = True 48 | self.cmd = Commands() 49 | # Open the connection to the db. We need to make this here. 50 | self.db = __database__ 51 | # When we exit, close the db 52 | atexit.register(self.db.close) 53 | self.prefix = '' 54 | 55 | def parse(self, data): 56 | root = '' 57 | args = [] 58 | 59 | # Split words by white space. 60 | words = data.split() 61 | # First word is the root command. 62 | root = words[0] 63 | 64 | # If there are more words, populate the arguments list. 65 | if len(words) > 1: 66 | args = words[1:] 67 | 68 | return (root, args) 69 | 70 | 71 | def print_output(self, output, filename): 72 | if not output: 73 | return 74 | if filename: 75 | with open(filename.strip(), 'a') as out: 76 | for entry in output: 77 | if entry['type'] == 'info': 78 | out.write('[*] {0}\n'.format(entry['data'])) 79 | elif entry['type'] == 'item': 80 | out.write(' [-] {0}\n'.format(entry['data'])) 81 | elif entry['type'] == 'warning': 82 | out.write('[!] {0}\n'.format(entry['data'])) 83 | elif entry['type'] == 'error': 84 | out.write('[!] {0}\n'.format(entry['data'])) 85 | elif entry['type'] == 'success': 86 | out.write('[+] {0}\n'.format(entry['data'])) 87 | elif entry['type'] == 'table': 88 | out.write(str(table( 89 | header=entry['data']['header'], 90 | rows=entry['data']['rows'] 91 | ))) 92 | out.write('\n') 93 | else: 94 | out.write('{0}\n'.format(entry['data'])) 95 | print_success("Output written to {0}".format(filename)) 96 | else: 97 | for entry in output: 98 | if entry['type'] == 'info': 99 | print_info(entry['data']) 100 | elif entry['type'] == 'item': 101 | print_item(entry['data']) 102 | elif entry['type'] == 'warning': 103 | print_warning(entry['data']) 104 | elif entry['type'] == 'error': 105 | print_error(entry['data']) 106 | elif entry['type'] == 'success': 107 | print_success(entry['data']) 108 | elif entry['type'] == 'table': 109 | print(table( 110 | header=entry['data']['header'], 111 | rows=entry['data']['rows'] 112 | )) 113 | else: 114 | print(entry['data']) 115 | 116 | def stop(self): 117 | # Stop main loop. 118 | self.active = False 119 | # Close the db 120 | print_info('Wait until the database is synced...') 121 | self.db.close() 122 | 123 | def create_folders(self): 124 | """ Create the folders for the program""" 125 | # The name of the folder should read from the configuration file 126 | home_folder = '~/.stf/' 127 | stf_home_folder = os.path.expanduser(home_folder) 128 | 129 | # Create the ~/.stf/ folder for storing the history and the database 130 | if os.path.exists(stf_home_folder) == False: 131 | os.makedirs(stf_home_folder) 132 | 133 | # if there is an history file, read from it and load the history 134 | # so that they can be loaded in the shell. 135 | # just store it in the home directory. 136 | self.history_path = os.path.expanduser(stf_home_folder+'.stfhistory') 137 | 138 | def start(self): 139 | from stf.core.dataset import __datasets__ 140 | # Logo. 141 | logo() 142 | self.db.list() 143 | 144 | # Setup shell auto-complete. 145 | def complete(text, state): 146 | # Try to autocomplete modules. 147 | mods = [i for i in __modules__ if i.startswith(text)] 148 | if state < len(mods): 149 | return mods[state] 150 | 151 | # Try to autocomplete commands. 152 | cmds = [i for i in self.cmd.commands if i.startswith(text)] 153 | if state < len(cmds): 154 | return cmds[state] 155 | 156 | # Then autocomplete paths. 157 | if text.startswith("~"): 158 | text = "{0}{1}".format(os.getenv("HOME"), text[1:]) 159 | return (glob.glob(text+'*')+[None])[state] 160 | 161 | # Auto-complete on tabs. 162 | readline.set_completer_delims(' \t\n;') 163 | readline.parse_and_bind('tab: complete') 164 | readline.parse_and_bind('set editing-mode vi') 165 | readline.set_completer(complete) 166 | 167 | 168 | # Save commands in history file. 169 | def save_history(path): 170 | readline.write_history_file(path) 171 | 172 | if os.path.exists(self.history_path): 173 | readline.read_history_file(self.history_path) 174 | 175 | # Register the save history at program's exit. 176 | atexit.register(save_history, path=self.history_path) 177 | 178 | # Main loop. 179 | while self.active: 180 | if __datasets__.current: 181 | self.prefix = red(__datasets__.current.get_name() + ': ') 182 | else: 183 | self.prefix = '' 184 | prompt = self.prefix + cyan('stf > ', True) 185 | 186 | # Wait for input from the user. 187 | try: 188 | data = raw_input(prompt).strip() 189 | except KeyboardInterrupt: 190 | print("") 191 | # Terminate on EOF. 192 | except EOFError: 193 | self.stop() 194 | print("") 195 | continue 196 | # Parse the input if the user provided any. 197 | else: 198 | # Skip if the input is empty. 199 | if not data: 200 | continue 201 | 202 | # Check for output redirection 203 | filename = False 204 | 205 | # If there is a > in the string, we assume the user wants to output to file. 206 | # We erase this because it was interfering with our > filter 207 | #if '>' in data: 208 | # temp = data.split('>') 209 | # data = temp[0] 210 | # filename = temp[1] 211 | 212 | # If the input starts with an exclamation mark, we treat the 213 | # input as a bash command and execute it. 214 | if data.startswith('!'): 215 | os.system(data[1:]) 216 | continue 217 | 218 | # Try to split commands by ; so that you can sequence multiple 219 | # commands at once. 220 | split_commands = data.split(';') 221 | for split_command in split_commands: 222 | split_command = split_command.strip() 223 | if not split_command: 224 | continue 225 | 226 | # If it's an internal command, we parse the input and split it 227 | # between root command and arguments. 228 | root, args = self.parse(split_command) 229 | 230 | # Check if the command instructs to terminate. 231 | if root in ('exit', 'quit'): 232 | self.stop() 233 | continue 234 | 235 | try: 236 | # If the root command is part of the embedded commands list we 237 | # execute it. 238 | if root in self.cmd.commands: 239 | self.cmd.commands[root]['obj'](*args) 240 | 241 | # If the root command is part of loaded modules, we initialize 242 | # the module and execute it. 243 | elif root in __modules__: 244 | module = __modules__[root]['obj']() 245 | module.set_commandline(args) 246 | module.run() 247 | 248 | self.print_output(module.output, filename) 249 | del(module.output[:]) 250 | else: 251 | print("Command not recognized.") 252 | except KeyboardInterrupt: 253 | pass 254 | except Exception as e: 255 | print_error("The command {0} raised an exception:".format(bold(root))) 256 | traceback.print_exc() 257 | 258 | --------------------------------------------------------------------------------