Read config from stdin
This commit is contained in:
@@ -30,6 +30,7 @@ class Proxy:
|
||||
def __init__(self, internal_port, public_port, callback_blocked_update=None, filters=None, public_host="0.0.0.0", internal_host="127.0.0.1"):
|
||||
self.filter_map = {}
|
||||
self.filter_map_lock = Lock()
|
||||
self.update_config_lock = Lock()
|
||||
self.public_host = public_host
|
||||
self.public_port = public_port
|
||||
self.internal_host = internal_host
|
||||
@@ -37,35 +38,27 @@ class Proxy:
|
||||
self.filters = set(filters) if filters else set([])
|
||||
self.process = None
|
||||
self.callback_blocked_update = callback_blocked_update
|
||||
self.config_file_path = None
|
||||
while self.config_file_path is None:
|
||||
config_file_path = os.path.join("/tmp/" + token_urlsafe(16))
|
||||
if not os.path.exists(config_file_path):
|
||||
self.config_file_path = config_file_path
|
||||
|
||||
def start(self, in_pause=False):
|
||||
if not self.isactive():
|
||||
self.filter_map = self.compile_filters()
|
||||
filters_codes = list(self.filter_map.keys()) if not in_pause else []
|
||||
proxy_binary_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),"./proxy")
|
||||
try:
|
||||
self.__write_config(filters_codes)
|
||||
|
||||
self.process = subprocess.Popen(
|
||||
[ proxy_binary_path, str(self.public_host), str(self.public_port), str(self.internal_host), str(self.internal_port), self.config_file_path],
|
||||
stdout=subprocess.PIPE, universal_newlines=True
|
||||
)
|
||||
for stdout_line in iter(self.process.stdout.readline, ""):
|
||||
if stdout_line.startswith("BLOCKED"):
|
||||
regex_id = stdout_line.split()[1]
|
||||
with self.filter_map_lock:
|
||||
self.filter_map[regex_id].blocked+=1
|
||||
if self.callback_blocked_update: self.callback_blocked_update(self.filter_map[regex_id])
|
||||
self.process.stdout.close()
|
||||
return self.process.wait()
|
||||
finally:
|
||||
self.__delete_config()
|
||||
|
||||
|
||||
self.process = subprocess.Popen(
|
||||
[ proxy_binary_path, str(self.public_host), str(self.public_port), str(self.internal_host), str(self.internal_port)],
|
||||
stdout=subprocess.PIPE, stdin=subprocess.PIPE, universal_newlines=True
|
||||
)
|
||||
self.update_config(filters_codes, sendsignal=False)
|
||||
|
||||
for stdout_line in iter(self.process.stdout.readline, ""):
|
||||
if stdout_line.startswith("BLOCKED"):
|
||||
regex_id = stdout_line.split()[1]
|
||||
with self.filter_map_lock:
|
||||
self.filter_map[regex_id].blocked+=1
|
||||
if self.callback_blocked_update: self.callback_blocked_update(self.filter_map[regex_id])
|
||||
self.process.stdout.close()
|
||||
return self.process.wait()
|
||||
|
||||
def stop(self):
|
||||
if self.isactive():
|
||||
@@ -84,36 +77,31 @@ class Proxy:
|
||||
self.start(in_pause=in_pause)
|
||||
return status
|
||||
|
||||
def __write_config(self, filters_codes):
|
||||
with open(self.config_file_path,'w') as config_file:
|
||||
for line in filters_codes:
|
||||
config_file.write(line + '\n')
|
||||
|
||||
def __delete_config(self):
|
||||
if os.path.exists(self.config_file_path):
|
||||
os.remove(self.config_file_path)
|
||||
def update_config(self, filters_codes, sendsignal=True):
|
||||
with self.update_config_lock:
|
||||
if (self.isactive()):
|
||||
self.process.stdin.write(" ".join(filters_codes))
|
||||
self.process.stdin.write(" END ")
|
||||
self.process.stdin.flush()
|
||||
if sendsignal:
|
||||
self.process.send_signal(SIGUSR1)
|
||||
|
||||
|
||||
def reload(self):
|
||||
if self.isactive():
|
||||
with self.filter_map_lock:
|
||||
self.filter_map = self.compile_filters()
|
||||
filters_codes = list(self.filter_map.keys())
|
||||
self.__write_config(filters_codes)
|
||||
self.trigger_reload_config()
|
||||
self.update_config(filters_codes)
|
||||
|
||||
def isactive(self):
|
||||
if self.process and not self.process.poll() is None:
|
||||
self.process = None
|
||||
return True if self.process else False
|
||||
|
||||
def trigger_reload_config(self):
|
||||
self.process.send_signal(SIGUSR1)
|
||||
|
||||
|
||||
def pause(self):
|
||||
if self.isactive():
|
||||
self.__write_config([])
|
||||
self.trigger_reload_config()
|
||||
self.update_config([])
|
||||
else:
|
||||
self.start(in_pause=True)
|
||||
|
||||
|
||||
@@ -73,13 +73,13 @@ struct regex_rules{
|
||||
if (arg[0] == '1'){
|
||||
regex regex(expr);
|
||||
#ifdef DEBUG
|
||||
cerr << "Added case sensitive regex " << expr_str << endl;
|
||||
cerr << "Added case sensitive regex " << expr << endl;
|
||||
#endif
|
||||
getByCode(arg[1])->push_back(make_pair(string(arg), regex));
|
||||
} else {
|
||||
regex regex(expr,regex_constants::icase);
|
||||
#ifdef DEBUG
|
||||
cerr << "Added case insensitive regex " << expr_str << endl;
|
||||
cerr << "Added case insensitive regex " << expr << endl;
|
||||
#endif
|
||||
getByCode(arg[1])->push_back(make_pair(string(arg), regex));
|
||||
}
|
||||
@@ -91,7 +91,6 @@ struct regex_rules{
|
||||
};
|
||||
shared_ptr<regex_rules> regex_config;
|
||||
|
||||
const char* config_file;
|
||||
mutex update_mutex;
|
||||
#ifdef MULTI_THREAD
|
||||
mutex stdout_mutex;
|
||||
@@ -393,25 +392,23 @@ namespace tcp_proxy
|
||||
|
||||
|
||||
void update_regex(){
|
||||
#ifdef DEBUG
|
||||
cerr << "Updating configuration" << endl;
|
||||
#endif
|
||||
std::unique_lock<std::mutex> lck(update_mutex);
|
||||
fstream fd;
|
||||
fd.open(config_file,ios::in);
|
||||
if (!fd.is_open()){
|
||||
cerr << "Error: config file couln't be opened" << endl;
|
||||
exit(1);
|
||||
}
|
||||
regex_rules *regex_new_config = new regex_rules();
|
||||
string line;
|
||||
while(getline(fd, line)) regex_new_config->add(line.c_str());
|
||||
string data;
|
||||
while(true){
|
||||
cin >> data;
|
||||
if (data == "END") break;
|
||||
regex_new_config->add(data.c_str());
|
||||
}
|
||||
regex_config.reset(regex_new_config);
|
||||
}
|
||||
|
||||
void signal_handler(int signal_num)
|
||||
{
|
||||
if (signal_num == SIGUSR1){
|
||||
#ifdef DEBUG
|
||||
cerr << "Updating configurtation" << endl;
|
||||
#endif
|
||||
update_regex();
|
||||
}else if(signal_num == SIGTERM){
|
||||
if (ios_loop != nullptr) ios_loop->stop();
|
||||
@@ -424,9 +421,9 @@ void signal_handler(int signal_num)
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
if (argc < 6)
|
||||
if (argc < 5)
|
||||
{
|
||||
cerr << "usage: tcpproxy_server <local host ip> <local port> <forward host ip> <forward port> <config_file>" << endl;
|
||||
cerr << "usage: tcpproxy_server <local host ip> <local port> <forward host ip> <forward port>" << endl;
|
||||
return 1;
|
||||
}
|
||||
|
||||
@@ -434,16 +431,19 @@ int main(int argc, char* argv[])
|
||||
const unsigned short forward_port = static_cast<unsigned short>(::atoi(argv[4]));
|
||||
const string local_host = argv[1];
|
||||
const string forward_host = argv[3];
|
||||
|
||||
config_file = argv[5];
|
||||
|
||||
update_regex();
|
||||
signal(SIGUSR1, signal_handler);
|
||||
signal(SIGTERM, signal_handler);
|
||||
|
||||
signal(SIGUSR1, signal_handler);
|
||||
|
||||
boost::asio::io_service ios;
|
||||
ios_loop = &ios;
|
||||
|
||||
|
||||
signal(SIGTERM, signal_handler);
|
||||
|
||||
#ifdef DEBUG
|
||||
cerr << "Starting Proxy" << endl;
|
||||
#endif
|
||||
try
|
||||
{
|
||||
tcp_proxy::bridge::acceptor acceptor(ios,
|
||||
@@ -470,6 +470,9 @@ int main(int argc, char* argv[])
|
||||
cerr << "Error: " << e.what() << endl;
|
||||
return 1;
|
||||
}
|
||||
#ifdef DEBUG
|
||||
cerr << "Proxy stopped!" << endl;
|
||||
#endif
|
||||
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user