#!/usr/bin/python3

# Copyright (C) 2022 - 2025
# Alexander Burmatov
# All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#   
# * Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.  
# * Redistributions in binary form must reproduce the above copyright notice, 
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.  
# * Neither the name of the Alexander Burmatov may be used to
#   endorse or promote products derived from this software without
#   specific prior written permission. 
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

'''
@author: Alexander Burmatov, thatman at altlinux dot org
'''

import argparse
import socket
import os
import sys
import secrets
import MySQLdb
import datetime
import uuid

sys.path.append('/etc/openuds/')
from settings import DATABASES


ip_addr = socket.gethostbyname(socket.gethostname())
creation_datetime = datetime.datetime.today()
version = '4.0.0'
uuid = str(uuid.uuid4())
subtype = 'Tunnel'

parser = argparse.ArgumentParser(description='Register tunnel token in MySQL DB')
parser.add_argument(
    '-H',
    '--host',
    type=str,
    default='',
    help='Input tunnel server IP Address'
)
parser.add_argument(
    '-p',
    '--port',
    type=int,
    default=10443,
    help='Input tunnel server port (default: 10443)'
)
parser.add_argument(
    '-n',
    '--name', 
    type=str, 
    default='',
    help='Input tunnel server name'
)
parser.add_argument(
    '-t',
    '--token',
    type=str,
    default='',
    help='Input tunnel server token (default: "")'
)
parser.add_argument(
    '-m',
    '--mac',
    type=str,
    default='',
    help='Input tunnel server MAC (default: "")'
)
parser.add_argument(
    '-c',
    '--certificate',
    type=str,
    default='---BEGIN CERT---...',
    help='Input tunnel server certificate (default: "---BEGIN CERT---...")'
)
parser.add_argument(
    '-l',
    '--log-level',
    type=int,
    default=3,
    help='Input tunnel server log level (0 = DEBUG, 1 = INFO, 2 = WARNING, 3 = ERROR) (default: 3)'
)
parser.add_argument(
    '-N',
    '--generate_new_token',
    type=bool,
    default=False,
    help='Input True if you want to generate a new token (default: False)'
)
args = parser.parse_args()

empty_name = args.name == ''
empty_ip = args.host == ''
empty_mac = args.mac == ''
only_token = args.token != '' and not args.generate_new_token
only_gen_new_token = args.token == '' and args.generate_new_token
if empty_ip:
    print('Empty tunnel server IP Address')
elif empty_name:
    print('Empty tunnel server name')
elif empty_mac:
    print('Empty tunnel server MAC')
elif args.token == '' and not args.generate_new_token:
    print('Choose to generate a new token or enter a token')
elif only_token != only_gen_new_token:
    if only_gen_new_token:
        token = secrets.token_urlsafe(36)
    else:
        token = args.token
    db=MySQLdb.connect(host=DATABASES['default']['HOST'], user=DATABASES['default']['USER'],
    passwd=DATABASES['default']['PASSWORD'], db=DATABASES['default']['NAME'])  
    c=db.cursor()
    c.execute("""INSERT INTO uds_server (register_username, register_ip, ip, hostname, token, stamp, uuid, type, subtype, maintenance_mode, version, data, os_type, mac, certificate, listen_port, log_level, locked_until) VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s);""", 
    (os.getlogin(), ip_addr, args.host, args.name, token, creation_datetime, uuid, 1, subtype, 0, version, None, 'Linux', args.mac, args.certificate, args.port, args.log_level, None,))
    db.commit()
    c.close()
    print(f'Tunnel token register success. (With token: {token})')
else:
    print('Choose to generate a new token only or only enter the token')