#!/usr/bin/perl

use strict;
use warnings;

use English;
use Getopt::Long;
use Path::Tiny;
use File::Temp;
use FindBin;
require "syscall.ph";

my $tmp_dir = File::Temp::tempdir(CLEANUP => 1);

*CLONE_NEWUTS  = \0x4000000;
*CLONE_NEWUSER = \0x10000000;
*CLONE_NEWNET  = \0x40000000;
our ( $CLONE_NEWUTS, $CLONE_NEWUSER, $CLONE_NEWNET );

sub exit_error {
    print STDERR "Error: ", $_[0], "\n";
    exit (exists $_[1] ? $_[1] : 1);
}

sub get_guidmapcmd {
    my ($guid) = (@_);
    my $config_file = "/etc/sub${guid}";
    my ($id) = ${guid} eq 'uid' ? ($UID) : split(' ', $GID);
    my ($current_usergroup) = ${guid} eq 'uid' ? getpwuid($id) : getgrgid($id);
    for my $line (path($config_file)->lines) {
        chomp $line;
        my ($usergroup, $lowerid, $count) = split(':', $line);
        next unless ($usergroup eq $current_usergroup || $usergroup eq $id);
        return "0 $id 1 1 $lowerid $count";
    }
    exit_error "Could not find $guid in $config_file";
}

sub start_namespace_socat {
    my ($opts) = @_;
    return () if $opts->{mode}{slirp4netns};

    my $socat_pid = fork() // exit_error("fork() failed: $!");
    if ($socat_pid == 0) {
        exec('socat', 'TCP-LISTEN:9050,reuseaddr,fork', "UNIX-CONNECT:$tmp_dir/torsocks.sock");
    }
    return ($socat_pid) unless $opts->{DNSPort};
    my $socat_pid_dnsport = fork() // exit_error("fork() failed: $!");
    if ($socat_pid_dnsport == 0) {
        exec('socat', 'UDP4-RECVFROM:53,reuseaddr,reuseport,fork', "UNIX-CONNECT:$tmp_dir/DNSPort.sock");
    }
    return ($socat_pid, $socat_pid_dnsport);
}

# starts the socat we run outside the namespace
sub start_socat {
    my ($opts) = @_;
    return () if $opts->{mode}{slirp4netns};

    my $socat_pid = fork() // exit_error("fork() failed: $!");
    if ($socat_pid == 0) {
        exec('socat', "UNIX-LISTEN:$tmp_dir/torsocks.sock,fork,reuseaddr,unlink-early,mode=700", "TCP:127.0.0.1:$opts->{SocksPort}");
    }
    return ($socat_pid) unless $opts->{DNSPort};
    my $socat_pid_dnsport = fork() // exit_error("fork() failed: $!");
    if ($socat_pid_dnsport == 0) {
        exec('socat', "UNIX-LISTEN:$tmp_dir/DNSPort.sock,fork,reuseaddr,unlink-early,mode=700", "UDP-SENDTO:127.0.0.1:$opts->{DNSPort}");
    }
    return ($socat_pid, $socat_pid_dnsport);
}

sub start_redsocks {
    my ($opts) = @_;
    pipe my $rfh, my $wfh;
    my $pid = fork() // exit_error("fork() failed: $!");
    if ($pid == 0) {
        close *STDOUT;
        close *STDERR;
        exec('/usr/sbin/redsocks', '-c', "$FindBin::Bin/redsocks.conf");
    }

    system(qw|ip route add local default dev lo|);
    system(qw|/usr/sbin/iptables -t nat -A OUTPUT -d 127.0.0.0/8 -j RETURN|);
    system(qw|/usr/sbin/iptables -t nat -A OUTPUT -p tcp --syn -j REDIRECT --to-ports 12345|);
    system(qw|/usr/sbin/iptables -t nat -A PREROUTING -p udp -m udp --dport 53 -j DNAT --to 127.0.0.1:53|)
        if $opts->{DNSPort};

    # TODO: When DNSPort is not used, we have no DNS working. Maybe we
    # can do something to fix that.

    return $pid;
}

sub setup_slirp4netns_firewall {
    my ($opts) = @_;

    # TODO: we need to add some iptable rules to redirect all DNS requests
    # to 127.0.0.1:53 and all tcp connections to 127.0.0.1:9052
    # and block all non-tor connections
}

sub unshare_run {
    my ($s, $opts) = @_;
    pipe my $rfh_slirp4netns, my $wfh_slirp4netns if $opts->{mode}{slirp4netns};
    my $f1 = fork();
    if ($f1 == 0) {
        my $ppid = $$;
        pipe my $rfh, my $wfh;
        my $pid = fork() // exit_error("fork() failed: $!");
        if ($pid == 0) {
            close $wfh;
            exit_error("read() did not receive EOF")
              unless sysread($rfh, my $c, 1) == 0;
            my $uidmapcmd = get_guidmapcmd('uid');
            exit_error("newuidmap $ppid $uidmapcmd failed: $!")
              unless system("newuidmap $ppid $uidmapcmd") == 0;
            my $gidmapcmd = get_guidmapcmd('gid');
            exit_error("newgidmap $ppid $gidmapcmd failed: $!")
              unless system("newgidmap $ppid $gidmapcmd") == 0;
            exit 0;
        }

        my $unshare_flags = $CLONE_NEWUSER | $CLONE_NEWUTS | $CLONE_NEWNET;
        syscall &SYS_unshare, $unshare_flags;
        close $wfh;
        waitpid($pid, 0) or exit_error("waitpid() failed: $!");
        exit_error("failed to set uidmap") if $? >> 8;
        syscall(&SYS_setgid, 0) == 0 or exit_error("setgid failed: $!");
        syscall(&SYS_setuid, 0) == 0 or exit_error("setuid failed: $!");
        syscall(&SYS_setgroups, 0, 0) == 0 or exit_error("setgroups failed: $!");

        my $f2 = fork() // exit_error("fork() failed: $!");
        if ($f2) {
            waitpid($f2, 0) or exit_error("waitpid() failed: $!");
            exit $? >> 8;
        }

        system(qw/ip link set lo up/);
        my @pids = start_namespace_socat($opts);

        push @pids, start_redsocks($opts) if $opts->{mode}{redsocks};

        if ($opts->{mode}{slirp4netns}) {
            syswrite($wfh_slirp4netns, $PID);
            close $wfh_slirp4netns;
            setup_slirp4netns_firewall($opts);
        }

        my $res = $s->();

        foreach my $spid (@pids) {
            kill 15, $spid;
        }
        exit($res ? $res >> 8 : 0);
    }
    my $pid_slirp4netns;
    if ($opts->{mode}{slirp4netns}) {
        $pid_slirp4netns = fork() // exit_error("fork() failed: $!");
        if ($pid_slirp4netns == 0) {
            exit_error("read() received EOF")
              if sysread($rfh_slirp4netns, my $c, 10) == 0;
            exec('slirp4netns', '--configure', $c, 'tap0');
        }
    }
    waitpid($f1, 0) or exit_error("waitpid() failed: $!");
    my $res = $? >> 8;
    kill 15, $pid_slirp4netns if $pid_slirp4netns;
    return $res;
}

sub usage {
    print STDERR <<USAGE_EOF
torsocks-netns [OPTIONS] -- [TORSOCKS-OPTIONS] [COMMAND [ARG...]]

Options:
  --help
    Print this message.

  --mode=<torsocks|redsocks|slirp4netns>
    Default mode is torsocks.

  --SocksPort=<port>
    Set Tor Socks port (default: 9050).

  --TransPort=<port>
    Set Tor transparent proxy port (TransPort in torrc). When this is set,
    instead of using torsocks we use some iptable rules to redirect all
    connections to Tor. This requires setting the --DNSPort option too.
    Using this option automatically selects the slirp4netns mode.

  --DNSPort=<port>
    Set Tor DNSPort. This can only be used in redsocks and slirp4netns modes.

USAGE_EOF
    ;
    exit 0;
}

my @cmd = @ARGV;
my %opts = ( SocksPort => '9050' );
my @options = ('SocksPort=s', 'TransPort=s', 'DNSPort=s', 'mode=s', 'help!');
Getopt::Long::GetOptionsFromArray(\@cmd, \%opts, @options) || exit 1;
usage if (!@cmd || $opts{help});

if (!$opts{mode}) {
    if ($opts{TransPort} || $opts{DNSPort}) {
        $opts{mode}{'slirp4netns'} = 1;
    } else {
        $opts{mode}{torsocks} = 1;
    }
} else {
    my %known_modes = map { $_ => 1 } qw/torsocks redsocks slirp4netns/;
    exit_error "Unknown mode $opts{mode}" unless $known_modes{$opts{mode}};
    $opts{mode} = { $opts{mode} => 1 };
}

if ($opts{TransPort} || $opts{mode}{slirp4netns}) {
    exit_error "--TransPort can only be used in slirp4netns mode"
        unless $opts{mode}{slirp4netns};
    exit_error "--TransPort argument is missing" unless $opts{TransPort};
    exit_error "--DNSPort argument is missing" unless $opts{DNSPort};
}

if ($opts{DNSPort} && $opts{mode}{torsocks}) {
    exit_error "--DNSPort cannot be used in torsocks mode";
}

if ($opts{mode}{slirp4netns}) {
    print STDERR "Warning: slirp4netns mode is not currently working.\n",
                 "Connections are not going through Tor.\n";
}

@cmd = ('torsocks', @cmd) if $opts{mode}{torsocks};

my @socat_pid = start_socat(\%opts);
my $res = unshare_run(
    sub {
        return system(@cmd);
    },
    \%opts
);
foreach my $pid (@socat_pid) {
    kill 15, $pid;
}
exit $res;
