-
Notifications
You must be signed in to change notification settings - Fork 0
/
clab.py
executable file
·130 lines (106 loc) · 4.33 KB
/
clab.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
#!/usr/bin/env python3
import os
import aiodocker
import asyncio
import argparse
import logging
import sys
import traceback
from dns import ContainerDnsServer
from proxies import ContainerProxy
from configuration import Configuration
from container import ContainerManager, ContainerBuilder
from models import close_database, setup_database
from netfilter import ContainerFirewall, TCPEndPoint, IPTableRules
from multiprocessing import Process, Pool
BANNER = """
_ _
| | | |
___| | __ _| |__
/ __| | / _` | '_ \
| (__| |___| (_| | |_) |
\___|______\__,_|_.__/
"""
CONTACT_INFO = """
Darren Southern
@stillinsecure
https://github.com/stillinsecure/clab
"""
def setup_logging(instance):
pid = os.getpid()
logging._srcfile = None
logging.logThreads = 1
logging.logProcesses = 0
if not os.path.exists('logs'):
os.makedirs('logs')
logging.basicConfig(filename='logs/{}-clab.out'.format(pid),
format='[{}:{}] %(asctime)s %(levelname)-8s %(message)s'.format(instance, pid),
level=logging.INFO,
datefmt='%Y-%m-%d %H:%M:%S')
def start_firewall(instance):
instance = instance - 1
try:
setup_logging(instance)
loop = asyncio.get_event_loop()
config.firewall.proxy_port = config.firewall.proxy_port + instance
config.firewall.queue_num = instance
container_mgr = ContainerManager(config)
loop.run_until_complete(container_mgr.start())
firewall = ContainerFirewall(container_mgr, config.firewall.queue_num)
firewall.start(loop)
proxy_server = ContainerProxy(container_mgr, config)
loop.run_until_complete(proxy_server.start(loop))
loop.run_forever()
finally:
loop.run_until_complete(container_mgr.stop())
def start_dns_server(config):
setup_logging(0)
container_mgr = ContainerManager(config)
with ContainerDnsServer(container_mgr, config.network.domain) as dns_server:
dns_server.start()
if __name__ == '__main__':
parser = argparse.ArgumentParser('Run clab')
parser.add_argument('--run', action='store_true',
help='Runs the current deployment')
parser.add_argument('--create', action='store_true',
help='Create and run a new deployment')
parser.add_argument('--delete', action='store_true',
help='Delete the current deployment')
parser.add_argument('--view', action='store_true',
help='View current deployment')
parser.add_argument('--config', default='config.yaml',
help='yaml file used to configure a deployment')
parser.add_argument('--ignore', default='store_true',
help='Ignore configuration changes detection')
parser.add_argument('--update', action='store_true',
help='Images defined in the config file are appended to the current deployment')
args = parser.parse_args()
print(BANNER)
print(CONTACT_INFO)
# Sets up the database for Peewee ORM
setup_database()
config = Configuration()
config.open(args.config)
# Displays the current deployment
if args.view:
builder = ContainerBuilder(config)
builder.view()
# Deletes the docker containers and entries in the containers.db
elif args.delete:
loop = asyncio.get_event_loop()
builder = ContainerBuilder(config)
loop.run_until_complete(builder.delete_containers())
elif args.create or args.run or args.update:
# Create a new deployment, this overwrites any existing deployment
if args.create or args.update:
setup_logging(0)
loop = asyncio.get_event_loop()
builder = ContainerBuilder(config)
loop.run_until_complete(builder.create_containers(args.update))
# Start the firewall/proxy using a worker pool
if args.run:
dns_server_process = Process(target=start_dns_server, args=(config, ))
dns_server_process.daemon = True
dns_server_process.start()
pool = Pool(config.firewall.instances)
pool.map(start_firewall, range(1, config.firewall.instances + 1))