├── .gitignore ├── .scrutinizer.yml ├── LICENSE ├── README.md ├── Vagrantfile ├── requirements.txt ├── sample_config ├── setup.py ├── tcpstat ├── __init__.py └── tcpstat.py ├── tests └── tests.py └── vagrant ├── bootstrap.sh └── sources.list /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | 5 | # C extensions 6 | *.so 7 | 8 | # Distribution / packaging 9 | .Python 10 | env/ 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | lib/ 17 | lib64/ 18 | parts/ 19 | sdist/ 20 | var/ 21 | *.egg-info/ 22 | .installed.cfg 23 | *.egg 24 | 25 | # PyInstaller 26 | # Usually these files are written by a python script from a template 27 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 28 | *.manifest 29 | *.spec 30 | 31 | # Installer logs 32 | pip-log.txt 33 | pip-delete-this-directory.txt 34 | 35 | # Unit test / coverage reports 36 | htmlcov/ 37 | .tox/ 38 | .coverage 39 | .cache 40 | nosetests.xml 41 | coverage.xml 42 | 43 | # Translations 44 | *.mo 45 | *.pot 46 | 47 | # Django stuff: 48 | *.log 49 | 50 | # Sphinx documentation 51 | docs/_build/ 52 | 53 | # PyBuilder 54 | target/ 55 | 56 | # Vagrant 57 | .vagrant/ 58 | -------------------------------------------------------------------------------- /.scrutinizer.yml: -------------------------------------------------------------------------------- 1 | checks: 2 | python: 3 | code_rating: true 4 | duplicate_code: true 5 | format_bad_indentation: 6 | indentation: '4 spaces' 7 | format_mixed_indentation: true 8 | format_line_too_long: 9 | max_length: '79' 10 | imports_relative_import: true 11 | imports_wildcard_import: true 12 | format_bad_whitespace: true 13 | format_multiple_statements: true 14 | basic_invalid_name: 15 | functions: '[a-z_][a-z0-9_]{2,30}$' 16 | variables: '[a-z_][a-z0-9_]{2,30}$' 17 | whitelisted_names: 'i,j,k,ex,Run,_' 18 | constants: '(([A-Z_][A-Z0-9_]*)|(__.*__))$' 19 | attributes: '[a-z_][a-z0-9_]{2,30}$' 20 | arguments: '[a-z_][a-z0-9_]{2,30}$' 21 | class_attributes: '([A-Za-z_][A-Za-z0-9_]{2,30}|(__.*__))$' 22 | inline_vars: '[A-Za-z_][A-Za-z0-9_]*$' 23 | classes: '[A-Z_][a-zA-Z0-9]+$' 24 | modules: '(([a-z_][a-z0-9_]*)|([A-Z][a-zA-Z0-9]+))$' 25 | methods: '[a-z_][a-z0-9_]{2,30}$' 26 | classes_no_self_argument: true 27 | classes_bad_mcs_method_argument: true 28 | classes_bad_classmethod_argument: true 29 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2014 Ivan Cai 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Tcpstat 2 | 3 | A TCP port traffic monitor written in Python. 4 | 5 | [![Latest Version](https://img.shields.io/pypi/v/tcpstat.svg?label=pypi)](https://pypi.python.org/pypi/tcpstat/) 6 | [![Scrutinizer Code Quality](https://scrutinizer-ci.com/g/caizixian/tcpstat/badges/quality-score.png?b=master)](https://scrutinizer-ci.com/g/caizixian/tcpstat/?branch=master) 7 | 8 | ## Dependencies: 9 | 10 | * MongoDB: an open-source document database. 11 | * PyMongo: a native Python driver for MongoDB. 12 | * python-iptables: Python bindings for iptables. 13 | 14 | ## Install dependencies: 15 | 16 | ``` 17 | sudo apt-get update 18 | sudo apt-get install -y mongodb python-pip python-dev build-essential 19 | sudo pip install -r /vagrant/requirements.txt 20 | ``` 21 | 22 | [How to secure my MongoDB?][3] 23 | 24 | ## Developing: 25 | 26 | * Download Vagrant box at https://cloud-images.ubuntu.com/vagrant/trusty/current/trusty-server-cloudimg-i386-vagrant-disk1.box 27 | * `vagrant box add ubuntu/trusty32 trusty-server-cloudimg-i386-vagrant-disk1.box` 28 | * `vagrant up` 29 | * `vagrant provision` 30 | 31 | ## Install 32 | 33 | This project is under heavy development. It is subject to major, breaking changes. 34 | 35 | Don't use it on production server or even testing server. 36 | 37 | ## Wiki 38 | 39 | [GitHub wiki page][1] 40 | 41 | ## License 42 | 43 | MIT 44 | 45 | ## Bugs and Issues 46 | 47 | * Feel free to create issue at [issue tracker][2] 48 | * And please feel free to make pull requests. 49 | 50 | [1]:https://github.com/caizixian/tcpstat/wiki 51 | [2]:https://github.com/caizixian/tcpstat/issues 52 | [3]:https://docs.mongodb.org/manual/administration/security/ 53 | -------------------------------------------------------------------------------- /Vagrantfile: -------------------------------------------------------------------------------- 1 | # -*- mode: ruby -*- 2 | # vi: set ft=ruby : 3 | 4 | # Vagrantfile API/syntax version. Don't touch unless you know what you're doing! 5 | VAGRANTFILE_API_VERSION = "2" 6 | 7 | Vagrant.configure(VAGRANTFILE_API_VERSION) do |config| 8 | # All Vagrant configuration is done here. The most common configuration 9 | # options are documented and commented below. For a complete reference, 10 | # please see the online documentation at vagrantup.com. 11 | 12 | # Every Vagrant virtual environment requires a box to build off of. 13 | config.vm.box = "ubuntu/trusty32" 14 | 15 | # The url from where the 'config.vm.box' box will be fetched if it 16 | # doesn't already exist on the user's system. 17 | config.vm.box_url = "https://cloud-images.ubuntu.com/vagrant/trusty/current/trusty-server-cloudimg-i386-vagrant-disk1.box" 18 | 19 | config.vm.provision :shell, path: "vagrant/bootstrap.sh" 20 | # Create a forwarded port mapping which allows access to a specific port 21 | # within the machine from a port on the host machine. In the example below, 22 | # accessing "localhost:8080" will access port 80 on the guest machine. 23 | # config.vm.network :forwarded_port, guest: 80, host: 8080 24 | 25 | # Create a private network, which allows host-only access to the machine 26 | # using a specific IP. 27 | config.vm.network :private_network, ip: "192.168.33.10" 28 | 29 | # Create a public network, which generally matched to bridged network. 30 | # Bridged networks make the machine appear as another physical device on 31 | # your network. 32 | config.vm.network :public_network 33 | 34 | # If true, then any SSH connections made will enable agent forwarding. 35 | # Default value: false 36 | # config.ssh.forward_agent = true 37 | 38 | # Share an additional folder to the guest VM. The first argument is 39 | # the path on the host to the actual folder. The second argument is 40 | # the path on the guest to mount the folder. And the optional third 41 | # argument is a set of non-required options. 42 | config.vm.synced_folder ".", "/vagrant" 43 | 44 | # Provider-specific configuration so you can fine-tune various 45 | # backing providers for Vagrant. These expose provider-specific options. 46 | # Example for VirtualBox: 47 | # 48 | # config.vm.provider :virtualbox do |vb| 49 | # # Don't boot with headless mode 50 | # vb.gui = true 51 | # 52 | # # Use VBoxManage to customize the VM. For example to change memory: 53 | # vb.customize ["modifyvm", :id, "--memory", "1024"] 54 | # end 55 | # 56 | # View the documentation for the provider you're using for more 57 | # information on available options. 58 | 59 | # Enable provisioning with Puppet stand alone. Puppet manifests 60 | # are contained in a directory path relative to this Vagrantfile. 61 | # You will need to create the manifests directory and a manifest in 62 | # the file hashcorp/trusty32.pp in the manifests_path directory. 63 | # 64 | # An example Puppet manifest to provision the message of the day: 65 | # 66 | # # group { "puppet": 67 | # # ensure => "present", 68 | # # } 69 | # # 70 | # # File { owner => 0, group => 0, mode => 0644 } 71 | # # 72 | # # file { '/etc/motd': 73 | # # content => "Welcome to your Vagrant-built virtual machine! 74 | # # Managed by Puppet.\n" 75 | # # } 76 | # 77 | # config.vm.provision :puppet do |puppet| 78 | # puppet.manifests_path = "manifests" 79 | # puppet.manifest_file = "site.pp" 80 | # end 81 | 82 | # Enable provisioning with chef solo, specifying a cookbooks path, roles 83 | # path, and data_bags path (all relative to this Vagrantfile), and adding 84 | # some recipes and/or roles. 85 | # 86 | # config.vm.provision :chef_solo do |chef| 87 | # chef.cookbooks_path = "../my-recipes/cookbooks" 88 | # chef.roles_path = "../my-recipes/roles" 89 | # chef.data_bags_path = "../my-recipes/data_bags" 90 | # chef.add_recipe "mysql" 91 | # chef.add_role "web" 92 | # 93 | # # You may also specify custom JSON attributes: 94 | # chef.json = { :mysql_password => "foo" } 95 | # end 96 | 97 | # Enable provisioning with chef server, specifying the chef server URL, 98 | # and the path to the validation key (relative to this Vagrantfile). 99 | # 100 | # The Opscode Platform uses HTTPS. Substitute your organization for 101 | # ORGNAME in the URL and validation key. 102 | # 103 | # If you have your own Chef Server, use the appropriate URL, which may be 104 | # HTTP instead of HTTPS depending on your configuration. Also change the 105 | # validation key to validation.pem. 106 | # 107 | # config.vm.provision :chef_client do |chef| 108 | # chef.chef_server_url = "https://api.opscode.com/organizations/ORGNAME" 109 | # chef.validation_key_path = "ORGNAME-validator.pem" 110 | # end 111 | # 112 | # If you're using the Opscode platform, your validator client is 113 | # ORGNAME-validator, replacing ORGNAME with your organization name. 114 | # 115 | # If you have your own Chef Server, the default validation client name is 116 | # chef-validator, unless you changed the configuration. 117 | # 118 | # chef.validation_client_name = "ORGNAME-validator" 119 | end 120 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | python-iptables 2 | pymongo 3 | requests 4 | -------------------------------------------------------------------------------- /sample_config: -------------------------------------------------------------------------------- 1 | [Groups] 2 | Name:Gp1,Gp2 3 | 4 | [Gp1] 5 | Port:80,8080-8081 6 | Webhook:http://localhost/api/v1/http_bandwidth 7 | 8 | [Gp2] 9 | Port:443 10 | Webhook:http://localhost/api/v1/https_bandwidth -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Always prefer setuptools over distutils 2 | from setuptools import setup, find_packages 3 | # from codecs import open # To use a consistent encoding 4 | from os import path 5 | here = path.abspath(path.dirname(__file__)) 6 | 7 | with open(path.join(here, 'README.md'), encoding='utf-8') as f: 8 | long_description = f.read() 9 | 10 | setup( 11 | name='tcpstat', 12 | 13 | # Versions should comply with PEP440. For a discussion on single-sourcing 14 | # the version across setup.py and the project code, see 15 | # https://packaging.python.org/en/latest/development.html#single-sourcing-the-version 16 | version='0.0.2', 17 | 18 | description='A TCP port traffic monitor.', 19 | long_description=long_description, 20 | 21 | # The project's main homepage. 22 | url='https://github.com/caizixian/tcpstat', 23 | 24 | # Author details 25 | author='Ivan Cai', 26 | author_email='caizixian@users.noreply.github.com', 27 | 28 | # Choose your license 29 | license='MIT', 30 | 31 | # See https://pypi.python.org/pypi?%3Aaction=list_classifiers 32 | classifiers=[ 33 | # How mature is this project? Common values are 34 | # 3 - Alpha 35 | # 4 - Beta 36 | # 5 - Production/Stable 37 | 'Development Status :: 3 - Alpha', 38 | 39 | # Indicate who your project is intended for 40 | 'Intended Audience :: System Administrators', 41 | 'Intended Audience :: Developers', 42 | 'Topic :: System :: Networking :: Monitoring', 43 | 44 | # Pick your license as you wish (should match "license" above) 45 | 'License :: OSI Approved :: MIT License', 46 | 47 | # Specify the Python versions you support here. In particular, ensure 48 | # that you indicate whether you support Python 2, Python 3 or both. 49 | 'Programming Language :: Python :: 2', 50 | 'Programming Language :: Python :: 2.6', 51 | 'Programming Language :: Python :: 2.7', 52 | 'Programming Language :: Python :: 3', 53 | 'Programming Language :: Python :: 3.3', 54 | 'Programming Language :: Python :: 3.4', 55 | 'Operating System :: POSIX :: Linux', 56 | ], 57 | 58 | # What does your project relate to? 59 | keywords='traffic accounting network', 60 | 61 | # You can just specify the packages manually here if your project is 62 | # simple. Or you can use find_packages(). 63 | packages=find_packages(exclude=['vagrant', 'docs', 'tests*']), 64 | 65 | # List run-time dependencies here. These will be installed by pip when your 66 | # project is installed. For an analysis of "install_requires" vs pip's 67 | # requirements files see: 68 | # https://packaging.python.org/en/latest/technical.html#install-requires-vs-requirements-files 69 | install_requires=['python-iptables', 'pymongo', 'requests'], 70 | 71 | # List additional groups of dependencies here (e.g. development dependencies). 72 | # You can install these using the following syntax, for example: 73 | # $ pip install -e .[dev,test] 74 | # extras_require = { 75 | # 'dev': ['check-manifest'], 76 | # 'test': ['coverage'], 77 | #}, 78 | 79 | # If there are data files included in your packages that need to be 80 | # installed, specify them here. If using Python 2.6 or less, then these 81 | # have to be included in MANIFEST.in as well. 82 | #package_data={ 83 | # 'tcpstat': ['README.md', 'LICENSE'] 84 | #}, 85 | 86 | # Although 'package_data' is the preferred approach, in some case you may 87 | # need to place data files outside of your packages. 88 | # see http://docs.python.org/3.4/distutils/setupscript.html#installing-additional-files 89 | # In this case, 'data_file' will be installed into '/my_data' 90 | #data_files=[('my_data', ['data/data_file'])], 91 | 92 | # To provide executable scripts, use entry points in preference to the 93 | # "scripts" keyword. Entry points provide cross-platform support and allow 94 | # pip to create the appropriate form of executable for the target platform. 95 | entry_points={ 96 | 'console_scripts': [ 97 | 'tcpstat=tcpstat:main', 98 | ], 99 | }, 100 | ) 101 | -------------------------------------------------------------------------------- /tcpstat/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/caizixian/tcpstat/91d33e5a62010b1e84f99d995a30c294b407d58c/tcpstat/__init__.py -------------------------------------------------------------------------------- /tcpstat/tcpstat.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | # -*- coding: utf-8 -*- 3 | 4 | # The MIT License (MIT) 5 | 6 | # Copyright (c) 2014 Ivan Cai 7 | 8 | # Permission is hereby granted, free of charge, to any person obtaining a copy 9 | # of this software and associated documentation files (the "Software"), to deal 10 | # in the Software without restriction, including without limitation the rights 11 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 12 | # copies of the Software, and to permit persons to whom the Software is 13 | # furnished to do so, subject to the following conditions: 14 | 15 | # The above copyright notice and this permission notice shall be included in all 16 | # copies or substantial portions of the Software. 17 | 18 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 19 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 20 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 21 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 22 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 23 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 24 | # SOFTWARE. 25 | 26 | from __future__ import absolute_import, division, print_function, with_statement 27 | import os 28 | import sys 29 | import argparse 30 | import logging 31 | import datetime 32 | import json 33 | 34 | __author__ = 'Ivan Cai' 35 | __version__ = '0.0.2' 36 | 37 | # The ConfigParser module has been renamed to configparser in Python 3. 38 | try: 39 | import configparser 40 | except ImportError: 41 | import ConfigParser as configparser 42 | 43 | try: 44 | import pymongo 45 | import iptc 46 | except ImportError: 47 | sys.exit("You don't have the required Python packages installed.") 48 | 49 | 50 | def get_version(): 51 | """Return a list which contains version number. Order is major, minor, micro.""" 52 | return __version__.split('.') 53 | 54 | 55 | def check_python(): 56 | """Check whether user's Python version meets our requirements.""" 57 | info = sys.version_info 58 | if info[0] == 2 and not info[1] >= 6: 59 | sys.exit("Python 2.6+ required") 60 | elif info[0] == 3 and not info[1] >= 3: 61 | sys.exit("Python 3.3+ required") 62 | elif info[0] not in [2, 3]: 63 | sys.exit("Python version not supported") 64 | 65 | 66 | def check_root(): 67 | """Check whether user run our script using root privilege.""" 68 | if not os.geteuid() == 0: 69 | sys.exit("You need to have root privileges to run this script.") 70 | 71 | 72 | def init(group_list): 73 | """Create init iptables rules""" 74 | with open("/etc/tcpstat.sh", "w") as iptables_init_script: 75 | # Set shebang 76 | iptables_init_script.write("#!/bin/bash\n") 77 | # Create new chain for security reason 78 | iptables_init_script.write("/sbin/iptables -N ACCT\n") 79 | # Flush existing rules in our custom chain 80 | iptables_init_script.write("/sbin/iptables -F ACCT\n") 81 | # Attach new chain 82 | iptables_init_script.write("/sbin/iptables -A FORWARD -j ACCT\n") 83 | iptables_init_script.write("/sbin/iptables -A INPUT -j ACCT\n") 84 | iptables_init_script.write("/sbin/iptables -A OUTPUT -j ACCT\n") 85 | for group in group_list: 86 | for port in group["Port"]: 87 | iptables_init_script.write( 88 | "/sbin/iptables -A ACCT -p tcp --dport " + str(port) + "\n") 89 | iptables_init_script.write( 90 | "/sbin/iptables -A ACCT -p tcp --sport " + str(port) + "\n") 91 | os.system("chmod +x /etc/tcpstat.sh") 92 | 93 | 94 | def find_config(args): 95 | """Find user's path of config file.""" 96 | if args.config == None and os.path.exists("/etc/tcpstat/config"): 97 | return "/etc/tcpstat/config" 98 | elif os.path.exists(args.config): 99 | return args.config 100 | else: 101 | return None 102 | 103 | 104 | def check_port_validity(port): 105 | """Check whether a port is valid.""" 106 | if 0 <= int(port) <= 65535: 107 | return True 108 | else: 109 | return False 110 | 111 | 112 | def read_config(path): 113 | """Check whether the path of config file is valid.""" 114 | if path == None: 115 | sys.exit("Config file doesn't exist.") 116 | else: 117 | config = configparser.ConfigParser() 118 | config.read(path) 119 | 120 | logging.info("Config file loaded.") 121 | 122 | groupname_list = config.get("Groups", "Name").split(",") 123 | group_list = [] 124 | for groups in groupname_list: 125 | # In config file 126 | # [Gp1] 127 | # Port:-1,2,65534-65537 128 | # Webhook:http://localhost/api/v1/tcpstats 129 | # To 130 | # {"Name": "Gp1", "Port": [2,65534,65535], "Webhook": "http://localhost/api/v1/tcpstats"} 131 | logging.debug("Loading Group " + groups) 132 | temp_dict = {"Name": groups} 133 | port_list = [] 134 | logging.debug(config.get(groups, "Port")) 135 | logging.debug(config.get(groups, "Port").split(",")) 136 | for port_str in config.get(groups, "Port").split(","): 137 | logging.debug("Catch a port str " + port_str) 138 | if '-' not in port_str: 139 | if port_str.isdigit() and check_port_validity(port_str): 140 | logging.debug("Appended a port " + port_str) 141 | port_list.append(int(port_str)) 142 | else: 143 | logging.error("You entered " + port_str + 144 | " which is not valid port number.") 145 | else: 146 | logging.debug("Catch a port range " + port_str) 147 | head = int(port_str.split('-')[0]) 148 | tail = int(port_str.split('-')[1]) 149 | map(port_list.append, 150 | filter(check_port_validity, range(head, tail + 1))) 151 | temp_dict.update({"Port": port_list}) 152 | temp_dict.update({"Webhook": config.get(groups, 'Webhook', '')}) 153 | group_list.append(temp_dict) 154 | logging.debug("Final list of groups") 155 | logging.debug(group_list) 156 | return group_list 157 | 158 | def check_migration_lock(): 159 | if os.path.exists("/var/lock/tcpstat.lock"): 160 | logging.error("Migration lock found. Quit.") 161 | sys.exit("Migration lock found. Quit.") 162 | 163 | def update_db(group_list): 164 | 165 | check_migration_lock() 166 | 167 | table = iptc.Table(iptc.Table.FILTER) 168 | client = pymongo.MongoClient('mongodb://localhost:27017/') 169 | db = client['tcpstat'] 170 | collection = db['accounting'] 171 | today_str = str(datetime.date.today()) 172 | 173 | logging.info("Connect to db") 174 | 175 | if not collection.find_one({"Name": group["Name"], "Time": today_str}): 176 | migrate_db(group_list) 177 | 178 | check_migration_lock() 179 | 180 | chain = iptc.Chain(table, 'ACCT') 181 | for group in group_list: 182 | for rule in chain.rules: 183 | for match in rule.matches: 184 | if not match.sport: 185 | port_number = str(match.dport) 186 | rule_type = "RX" 187 | else: 188 | port_number = str(match.sport) 189 | rule_type = "TX" 190 | entry = collection.find_one( 191 | {"Name": group["Name"], "Time": today_str}) 192 | if port_number in entry.keys(): 193 | # Counters in bytes 194 | if rule_type == "TX": 195 | # Fetch entries in db 196 | entry = collection.find_one( 197 | {"Name": group["Name"], "Time": today_str}) 198 | TX = rule.get_counters()[1] + entry[port_number]["TX"] 199 | RX = entry[port_number]["RX"] 200 | logging.debug( 201 | "This record is TX " + str(TX) + " for " + port_number) 202 | # Modify data in db 203 | collection.update({"Name": group["Name"], 204 | "Time": today_str}, 205 | {"$set": 206 | {port_number: { 207 | "TX": int(TX), "RX": int(RX)}} 208 | } 209 | ) 210 | else: 211 | entry = collection.find_one( 212 | {"Name": group["Name"], "Time": today_str}) 213 | # Fetch entries in db 214 | RX = rule.get_counters()[1] + entry[port_number]["RX"] 215 | TX = entry[port_number]["TX"] 216 | logging.debug( 217 | "This record is RX " + str(RX) + " for " + port_number) 218 | collection.update({"Name": group["Name"], 219 | "Time": today_str}, 220 | {"$set": 221 | {port_number: { 222 | "TX": int(TX), "RX": int(RX)}} 223 | } 224 | ) 225 | 226 | chain.zero_counters() 227 | 228 | 229 | def migrate_db(group_list): 230 | #Migration lock 231 | open('/var/lock/tcpstat.lock', 'a').close() 232 | client = pymongo.MongoClient('mongodb://localhost:27017/') 233 | db = client['tcpstat'] 234 | collection = db['accounting'] 235 | today_str = str(datetime.date.today()) 236 | for group in group_list: 237 | entry = collection.find_one({"Name": group["Name"], "Time": today_str}) 238 | if entry: 239 | logging.info("Find an existing entry. Let's migrate it.") 240 | for port in group["Port"]: 241 | if str(port) not in entry.keys(): 242 | collection.update({"Name": group["Name"], "Time": today_str}, 243 | {"$set": {str(port): {"TX": 0, "RX": 0}}}) 244 | 245 | else: 246 | logging.info("Create a new entry with new schema.") 247 | temp_dict = {} 248 | temp_dict.update({"Name": group["Name"], "Time": today_str}) 249 | for port in group["Port"]: 250 | temp_dict.update({str(port): {"TX": 0, "RX": 0}}) 251 | collection.insert(temp_dict) 252 | #Remove migration lock 253 | os.remove('/var/lock/tcpstat.lock') 254 | 255 | 256 | def main(): 257 | # Check whether user run our script using root privilege. 258 | check_root() 259 | # Check whether user's Python version meets our requirements. 260 | check_python() 261 | 262 | # Setting up logging module. 263 | logging.basicConfig(filename='/var/log/tcpstat.log', 264 | format='%(asctime)s %(levelname)s: %(message)s', 265 | level=logging.DEBUG) 266 | 267 | logging.info(" ".join(("Started. Version:", __version__))) 268 | 269 | # Init command line argument. 270 | parser = argparse.ArgumentParser() 271 | parser.add_argument("-c", "--config", type=str, 272 | help="Path of config file. Default /etc/tcpstat/config") 273 | group = parser.add_mutually_exclusive_group() 274 | group.add_argument("-v", "--version", help="Show version.", 275 | action="store_true") 276 | group.add_argument("-i", "--init", help="Init iptables rules.", 277 | action="store_true") 278 | group.add_argument("-u", "--update", help="Update db with latest data.", 279 | action="store_true") 280 | group.add_argument("-m", "--migrate", help="Migrate db with new config.", 281 | action="store_true") 282 | 283 | # Parse command line argument. 284 | args = parser.parse_args() 285 | 286 | if args.version: 287 | print(" ".join(("Tcpstat\nVersion:", __version__))) 288 | 289 | # Init rules in /etc/tcpstat.sh which will be included in /etc/rc.local 290 | if args.init: 291 | init(read_config(find_config(args))) 292 | 293 | if args.update: 294 | update_db(read_config(find_config(args))) 295 | 296 | if args.migrate: 297 | migrate_db(read_config(find_config(args))) 298 | 299 | logging.info("Exit.") 300 | 301 | if __name__ == "__main__": 302 | main() 303 | -------------------------------------------------------------------------------- /tests/tests.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | # -*- coding: utf-8 -*- 3 | 4 | # The MIT License (MIT) 5 | 6 | # Copyright (c) 2014 Ivan Cai 7 | 8 | # Permission is hereby granted, free of charge, to any person obtaining a copy 9 | # of this software and associated documentation files (the "Software"), to deal 10 | # in the Software without restriction, including without limitation the rights 11 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 12 | # copies of the Software, and to permit persons to whom the Software is 13 | # furnished to do so, subject to the following conditions: 14 | 15 | # The above copyright notice and this permission notice shall be included in all 16 | # copies or substantial portions of the Software. 17 | 18 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 19 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 20 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 21 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 22 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 23 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 24 | # SOFTWARE. 25 | 26 | import tcpstat -------------------------------------------------------------------------------- /vagrant/bootstrap.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | cp /vagrant/vagrant/sources.list /etc/apt/sources.list 3 | apt-get update 4 | apt-get install -y mongodb python-pip python-dev build-essential 5 | apt-get clean 6 | pip install -r /vagrant/requirements.txt -i http://pypi.douban.com/simple -------------------------------------------------------------------------------- /vagrant/sources.list: -------------------------------------------------------------------------------- 1 | ## Note, this file is written by cloud-init on first boot of an instance 2 | ## modifications made here will not survive a re-bundle. 3 | ## if you wish to make changes you can: 4 | ## a.) add 'apt_preserve_sources_list: true' to /etc/cloud/cloud.cfg 5 | ## or do the same in user-data 6 | ## b.) add sources in /etc/apt/sources.list.d 7 | ## c.) make changes to template file /etc/cloud/templates/sources.list.tmpl 8 | # 9 | 10 | # See http://help.ubuntu.com/community/UpgradeNotes for how to upgrade to 11 | # newer versions of the distribution. 12 | deb http://mirrors.ustc.edu.cn/ubuntu trusty main 13 | deb-src http://mirrors.ustc.edu.cn/ubuntu trusty main 14 | 15 | ## Major bug fix updates produced after the final release of the 16 | ## distribution. 17 | deb http://mirrors.ustc.edu.cn/ubuntu trusty-updates main 18 | deb-src http://mirrors.ustc.edu.cn/ubuntu trusty-updates main 19 | 20 | ## N.B. software from this repository is ENTIRELY UNSUPPORTED by the Ubuntu 21 | ## team. Also, please note that software in universe WILL NOT receive any 22 | ## review or updates from the Ubuntu security team. 23 | deb http://mirrors.ustc.edu.cn/ubuntu trusty universe 24 | deb-src http://mirrors.ustc.edu.cn/ubuntu trusty universe 25 | deb http://mirrors.ustc.edu.cn/ubuntu trusty-updates universe 26 | deb-src http://mirrors.ustc.edu.cn/ubuntu trusty-updates universe 27 | 28 | ## N.B. software from this repository is ENTIRELY UNSUPPORTED by the Ubuntu 29 | ## team, and may not be under a free licence. Please satisfy yourself as to 30 | ## your rights to use the software. Also, please note that software in 31 | ## multiverse WILL NOT receive any review or updates from the Ubuntu 32 | ## security team. 33 | deb http://mirrors.ustc.edu.cn/ubuntu trusty multiverse 34 | deb-src http://mirrors.ustc.edu.cn/ubuntu trusty multiverse 35 | deb http://mirrors.ustc.edu.cn/ubuntu trusty-updates multiverse 36 | deb-src http://mirrors.ustc.edu.cn/ubuntu trusty-updates multiverse 37 | 38 | ## Uncomment the following two lines to add software from the 'backports' 39 | ## repository. 40 | ## N.B. software from this repository may not have been tested as 41 | ## extensively as that contained in the main release, although it includes 42 | ## newer versions of some applications which may provide useful features. 43 | ## Also, please note that software in backports WILL NOT receive any review 44 | ## or updates from the Ubuntu security team. 45 | deb http://mirrors.ustc.edu.cn/ubuntu trusty-backports main restricted universe multiverse 46 | deb-src http://mirrors.ustc.edu.cn/ubuntu trusty-backports main restricted universe multiverse 47 | 48 | ## Uncomment the following two lines to add software from Canonical's 49 | ## 'partner' repository. 50 | ## This software is not part of Ubuntu, but is offered by Canonical and the 51 | ## respective vendors as a service to Ubuntu users. 52 | # deb http://archive.canonical.com/ubuntu trusty partner 53 | # deb-src http://archive.canonical.com/ubuntu trusty partner 54 | 55 | deb http://mirrors.ustc.edu.cn/ubuntu trusty-security main 56 | deb-src http://mirrors.ustc.edu.cn/ubuntu trusty-security main 57 | deb http://mirrors.ustc.edu.cn/ubuntu trusty-security universe 58 | deb-src http://mirrors.ustc.edu.cn/ubuntu trusty-security universe 59 | deb http://mirrors.ustc.edu.cn/ubuntu trusty-security multiverse 60 | deb-src http://mirrors.ustc.edu.cn/ubuntu trusty-security multiverse 61 | --------------------------------------------------------------------------------