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
|
#!/usr/bin/env python3
#
# Copyright (C) 2023 VyOS maintainers and contributors
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
# Control pmacct daemons in a tricky way.
# Pmacct has signal processing in a main loop, together with packet
# processing. Because of this, while it is waiting for packets, it cannot
# handle the control signal. We need to start the systemctl command and then
# send some packets to pmacct to wake it up
from argparse import ArgumentParser
from socket import socket, AF_INET, SOCK_DGRAM
from sys import exit
from time import sleep
from psutil import Process
def stop_process(pid: int, timeout: int) -> None:
"""Send a signal to uacctd
and then send packets to special address predefined in a firewall
to unlock main loop in uacctd and finish the process properly
Args:
pid (int): uacctd PID
timeout (int): seconds to wait for a process end
"""
# find a process
uacctd = Process(pid)
uacctd.terminate()
# create a socket
trigger = socket(AF_INET, SOCK_DGRAM)
first_cycle: bool = True
while uacctd.is_running() and timeout:
print('sending a packet to uacctd...')
trigger.sendto(b'WAKEUP', ('127.0.254.0', 1))
# do not sleep during first attempt
if not first_cycle:
sleep(1)
timeout -= 1
first_cycle = False
if __name__ == '__main__':
parser = ArgumentParser()
parser.add_argument('process_id',
type=int,
help='PID file of uacctd core process')
parser.add_argument('timeout',
type=int,
help='time to wait for process end')
args = parser.parse_args()
stop_process(args.process_id, args.timeout)
exit()
|
