#!/usr/local/bin/python3.11

##### BSD Cafe CheckMyIP Server #####
##### Originally written by John W Kerns #####
##### Forked to BSD Cafe by Stefano Marinelli #####
##### https://brew.bsd.cafe/BSDCafe/checkmyip #####

##### Inform version here #####
version = "v1.3.1"  # Updated version

##### Import necessary modules #####
import os
import sys
import time
import socket
import json
import jinja2
import paramiko
import threading
from concurrent.futures import ThreadPoolExecutor
from dnslib import DNSRecord, DNSHeader, RR, TXT, QTYPE
import logging  # Use built-in logging module

##### Configure Logging #####
LOG_DIR = "/var/log/checkmyip/"
os.makedirs(LOG_DIR, exist_ok=True)

# General Logger
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s: %(message)s",
    handlers=[
        logging.FileHandler(os.path.join(LOG_DIR, f"{time.strftime('%Y-%m-%d')}_checkmyip.log")),
        logging.StreamHandler(sys.stdout)
    ]
)

# Paramiko Logger
paramiko_log = os.path.join(LOG_DIR, f"{time.strftime('%Y-%m-%d')}_ssh.log")
paramiko.util.log_to_file(paramiko_log)

##### Jinja formatting for logging queries #####
j2log = "Connection from: {{ ip }} ({{ port }}) ({{ proto }})"

##### Jinja formatting for response queries #####
j2send = """{
"comment": "##     Your IP Address is {{ ip }} ({{ port }})     ##",
"family": "{{ family }}",
"ip": "{{ ip }}",
"port": "{{ port }}",
"protocol": "{{ proto }}",
"version": "%s",
"website": "https://brew.bsd.cafe/BSDCafe/checkmyip",
"sponsor": "Served by BSD Cafe, https://bsd.cafe/"
}""" % version

##### DNS Server Functionality #####
class DNSResponder:
    def __init__(self, address='::', port=5553):
        self.address = address
        self.port = port
        self.sock = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
        self.sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)  # Enables dual-stack mode
        self.sock.bind((self.address, self.port))

    def run(self):
        logging.info("DNS Server listening on port %s for both IPv4 and IPv6", self.port)
        while True:
            try:
                data, addr = self.sock.recvfrom(1024)
                threading.Thread(target=self.handle_query, args=(data, addr), daemon=True).start()
            except Exception as e:
                logging.error("DNS Server error: %s", e)

    def handle_query(self, data, addr):
        try:
            request = DNSRecord.parse(data)
            qname = request.q.qname
            qtype = request.q.qtype
            ip, port, *_ = addr  # Get IP and port; ignore other details

            # Normalize IPv6 address (remove IPv4 mapping if exists)
            if ip.startswith("::ffff:"):
                ip = ip[7:]

            # Only respond to TXT queries for myip.bsd.cafe
            if qname.matchGlob('myip.bsd.cafe') and qtype == QTYPE.TXT:
                response = DNSRecord(DNSHeader(id=request.header.id, qr=1, aa=1, ra=1), q=request.q)
                response.add_answer(RR(qname, QTYPE.TXT, ttl=60, rdata=TXT(ip)))
                self.sock.sendto(response.pack(), addr)
        except Exception as e:
            logging.error("Error handling DNS query: %s", e)

##### Creates an RSA key for use by paramiko #####
class RSAKey:
    data = """-----BEGIN RSA PRIVATE KEY-----
MIICWgIBAAKBgQDTj1bqB4WmayWNPB+8jVSYpZYk80Ujvj680pOTh2bORBjbIAyz
oWGW+GUjzKxTiiPvVmxFgx5wdsFvF03v34lEVVhMpouqPAYQ15N37K/ir5XY+9m/
d8ufMCkjeXsQkKqFbAlQcnWMCRnOoPHS3I4vi6hmnDDeeYTSRvfLbW0fhwIBIwKB
gBIiOqZYaoqbeD9OS9z2K9KR2atlTxGxOJPXiP4ESqP3NVScWNwyZ3NXHpyrJLa0
EbVtzsQhLn6rF+TzXnOlcipFvjsem3iYzCpuChfGQ6SovTcOjHV9z+hnpXvQ/fon
soVRZY65wKnF7IAoUwTmJS9opqgrN6kRgCd3DASAMd1bAkEA96SBVWFt/fJBNJ9H
tYnBKZGw0VeHOYmVYbvMSstssn8un+pQpUm9vlG/bp7Oxd/m+b9KWEh2xPfv6zqU
avNwHwJBANqzGZa/EpzF4J8pGti7oIAPUIDGMtfIcmqNXVMckrmzQ2vTfqtkEZsA
4rE1IERRyiJQx6EJsz21wJmGV9WJQ5kCQQDwkS0uXqVdFzgHO6S++tjmjYcxwr3g
H0CoFYSgbddOT6miqRskOQF3DZVkJT3kyuBgU2zKygz52ukQZMqxCb1fAkASvuTv
qfpH87Qq5kQhNKdbbwbmd2NxlNabazPijWuphGTdW0VfJdWfklyS2Kr+iqrs/5wV
HhathJt636Eg7oIjAkA8ht3MQ+XSl9yIJIS8gVpbPxSw5OMfw0PjVE7tBdQruiSc
nvuQES5C9BMHjF39LZiGH1iLQy7FgdHyoP+eodI7
-----END RSA PRIVATE KEY-----
"""
    def readlines(self):  # For use by paramiko.RSAKey.from_private_key
        return self.data.split("\n")

    def get_key(self):
        return paramiko.RSAKey.from_private_key(self)

##### Imports and modifies the ServerInterface module for use by paramiko #####
class SSHServer(paramiko.ServerInterface):
    def __init__(self):
        self.event = threading.Event()

    def check_channel_request(self, kind, chanid):
        if kind == 'session':
            return paramiko.OPEN_SUCCEEDED
        return paramiko.OPEN_FAILED_ADMINISTRATIVELY_PROHIBITED

    def check_auth_none(self, username):
        return paramiko.AUTH_SUCCESSFUL

    def get_allowed_auths(self, username):
        return 'none'

    def check_channel_shell_request(self, channel):
        self.event.set()
        return True

    def check_channel_pty_request(self, channel, term, width, height, pixelwidth, pixelheight, modes):
        return True

##### Method to merge Jinja templates #####
def j2format(j2tmp, valdict):
    template = jinja2.Template(j2tmp)
    return template.render(valdict).replace("\n", "\r\n")

##### Cleans IP addresses coming from socket library #####
def cleanip(addr):
    ip, port = addr[:2]
    family = "ipv6"  # Default to IPv6
    if ip.startswith("::ffff:"):  # Check if this is a prefixed IPv4 address
        ip = ip.replace("::ffff:", "")  # Clean the IP
        family = "ipv4"
    elif ":" not in ip:  # Simple check to recognize IPv4
        family = "ipv4"
    return ip, port, family  # Return cleaned IP and family

##### TCP listener methods using ThreadPoolExecutor for better thread management #####
def listener(port, talker, executor):
    listen_ip = ''
    listen_port = port
    buffer_size = 1024
    while True:
        try:
            # Use AF_INET6 but also handle IPv4 connections
            sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
            sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)  # Allow IPv4-mapped IPv6 addresses
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            sock.bind((listen_ip, listen_port))
            sock.listen(buffer_size)
            logging.info("Listening on port %s", listen_port)
            while True:
                client, addr = sock.accept()
                ip, port, family = cleanip(addr)  # Get cleaned IP info including family
                valdict = {"ip": ip, "port": port, "family": family, "proto": ""}
                executor.submit(talker, client, valdict)  # Submit to thread pool
        except Exception as e:
            logging.error("Listener on port %s encountered an error: %s", port, e)
            time.sleep(5)  # Wait before retrying

##### Telnet responder method. Is run in own thread for each telnet query #####
def telnet_talker(client, valdict, proto="telnet"):
    try:
        valdict.update({"proto": proto})  # Add the protocol to the value dict
        logging.info(j2format(j2log, valdict))  # Log the query
        # Send the query response
        response = j2format(j2send, valdict) + "\n"
        client.sendall(response.encode())
    except Exception as e:
        logging.error("Telnet talker error: %s", e)
    finally:
        client.close()  # Ensure the socket is closed

##### SSH responder method. Gets run in own thread for each SSH query #####
def ssh_talker(client, valdict, proto="ssh"):
    try:
        def makefile(chan):
            # A hack to make SSH sessions work properly
            chan.makefile('rU').readline().strip('\r\n')

        valdict.update({"proto": proto})
        logging.info(j2format(j2log, valdict))
        t = paramiko.Transport(client)
        t.set_gss_host(socket.getfqdn(""))
        t.load_server_moduli()
        t.add_server_key(RSAKey().get_key())  # RSA key object
        server = SSHServer()
        t.start_server(server=server)

        chan = t.accept(20)
        if chan:
            if not server.event.wait(10):
                logging.warning("No shell request received.")
                chan.close()
            else:
                response = j2format(j2send, valdict) + "\n"
                chan.send(response)
                threading.Thread(target=makefile, args=(chan,), daemon=True).start()
                time.sleep(1)
                chan.close()
    except Exception as e:
        logging.error("SSH talker error: %s", e)
    finally:
        client.close()

##### HTTP responder method. Gets run in own thread for each HTTP query #####
def http_talker(client, valdict, proto="http"):
    try:
        time.sleep(0.1)  # Allow client to send data
        client.settimeout(5.0)  # Set timeout to prevent blocking indefinitely
        browser = False  # Is the client using a browser?
        headers = {}
        try:
            data = client.recv(2048)  # Receive data from the buffer (if any)
            if data:
                request_lines = data.decode().split('\r\n')
                for line in request_lines:
                    if ": " in line:
                        key, value = line.split(": ", 1)
                        headers[key] = value
                # Extract real IP and adjust family based on headers if present
                if 'X-Real-IP' in headers or 'X-Forwarded-For' in headers:
                    real_ip = headers.get('X-Real-IP', headers.get('X-Forwarded-For').split(',')[0])
                    ip, port, family = cleanip((real_ip, valdict['port']))
                    valdict.update({"ip": ip, "family": family})
                browser = True  # Set client browser to True
        except socket.timeout:
            browser = False  # Set client browser to False
        except Exception as e:
            logging.error("HTTP talker recv error: %s", e)
            browser = False

        if not browser:  # If the client is not a browser
            telnet_talker(client, valdict, "http-telnet")  # Hand to telnet_talker
        else:  # If client is a browser
            valdict.update({"proto": proto})
            logging.info(j2format(j2log, valdict))
            response_body_raw = j2format(j2send, valdict) + "\n"
            if request_lines and "/raw" in request_lines[0]:
                try:
                    json_content = json.loads(response_body_raw)
                    response_body_raw = json_content.get('ip', '')
                except json.JSONDecodeError:
                    response_body_raw = "Invalid JSON response."
            response_headers_raw = (
                "HTTP/1.1 200 OK\r\n"
                f"Content-Length: {len(response_body_raw)}\r\n"
                "Content-Type: application/json; charset=utf-8\r\n"
                "Connection: close\r\n\r\n"
            )
            client.sendall(response_headers_raw.encode() + response_body_raw.encode())
    except Exception as e:
        logging.error("HTTP talker error: %s", e)
    finally:
        client.close()  # Ensure the socket is closed

##### Server startup method. Starts a listener thread for each TCP port #####
def start():
    talkers = {
        2222: ssh_talker,
        2223: telnet_talker,
        2280: http_talker
    }  # Three listeners on different ports
    max_workers = 100  # Define maximum number of threads
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        for port, talker in talkers.items():
            # Launch a thread for each listener
            threading.Thread(target=listener, args=(port, talker, executor), daemon=True).start()
        # Start DNS server in a new thread
        dns_server = DNSResponder(port=5553)  # Assuming DNS is on port 5553 as per original code
        threading.Thread(target=dns_server.run, daemon=True).start()
        logging.info("All services started.")
        while True:  # Keep the main thread alive
            try:
                time.sleep(1)
            except KeyboardInterrupt:
                logging.info("Shutting down server.")
                sys.exit(0)

##### Client class to be used to make API calls to CheckMyIP server #####
class CheckMyIP_Client:
    def __init__(self):
        self._json = json  # Direct reference to json module
        self._socket = socket  # Direct reference to socket module
        self._raw_data = None  # Initialize the _raw_data variable
        self._data = None  # Initialize the _data variable
        self._af = "auto"  # Set the IP address family type to "auto"
        self.server = "telnetmyip.com"  # Set the default CheckMyIP server

    def get(self):  # Primary method to run IP check
        sock = None
        try:
            if self._af == "auto":  # If we are using an auto address family
                try:  # Try using IPv6
                    sock = self._socket.socket(self._socket.AF_INET6, self._socket.SOCK_STREAM)
                    sock.connect((self.server, 23))
                except:  # Fall back to IPv4 if IPv6 fails
                    sock = self._socket.socket(self._socket.AF_INET, self._socket.SOCK_STREAM)
                    sock.connect((self.server, 23))
            elif self._af == "ipv6":  # If we are using the IPv6 address family
                sock = self._socket.socket(self._socket.AF_INET6, self._socket.SOCK_STREAM)
                sock.connect((self.server, 23))
            elif self._af == "ipv4":  # If we are using the IPv4 address family
                sock = self._socket.socket(self._socket.AF_INET, self._socket.SOCK_STREAM)
                sock.connect((self.server, 23))
            self._raw_data = sock.recv(1024).decode()
            self._data = self._json.loads(self._raw_data)  # Receive data from the buffer
            return self._data  # Return the JSON data
        except Exception as e:
            logging.error("CheckMyIP_Client error: %s", e)
            return None
        finally:
            if sock:
                sock.close()  # Ensure the socket is closed

    def set_family(self, family):  # Method to set the IP address family
        allowed = ["auto", "ipv4", "ipv6"]  # Allowed input values
        if family in allowed:
            self._af = family
        else:
            raise ValueError("Allowed families are 'auto', 'ipv4', 'ipv6'")

### CheckMyIP_Client Example Usage ###
# client = CheckMyIP_Client()
# data = client.get()
# print(data)

if __name__ == "__main__":
    start()  # Start the server