# -*- coding: binary -*-

module Msf
  module Exploit::Remote::SMB
    # This mixin provides a minimal SMB server
    module Server
      include ::Msf::Exploit::Remote::SocketServer
      include ::Msf::Exploit::Remote::SMB::LogAdapter

      module ServiceMixin
        def start
          if Rex::Socket.is_ipv6?(@socket.localhost)
            localinfo = "[#{@socket.localhost}]:#{@socket.localport}"
          else
            localinfo = "#{@socket.localhost}:#{@socket.localport}"
          end

          self.listener_thread = Rex::ThreadFactory.spawn("SMBServerListener(#{localinfo})", false) do
            begin
              run do |server_client|
                on_client_connect_proc.call(server_client) if on_client_connect_proc
                true
              end
            rescue IOError => e
              # this 'IOError: stream closed in another thread' is expected, so disregard it
              wlog("#{e.class}: #{e.message}")
            end
          end
        end

        def stop
          @socket.close
        end

        def wait
          listener_thread.join if listener_thread
        end

        attr_accessor :listener_thread, :on_client_connect_proc
      end

      def initialize(info = {})
        super

        register_options(
          [
            OptPort.new('SRVPORT', [ true, 'The local port to listen on.', 445 ])
          ], self.class)
      end

      def start_service(opts = {})
        @rsock = Rex::Socket::Tcp.create(
          'LocalHost' => bindhost,
          'LocalPort' => bindport,
          'Comm' => _determine_server_comm(bindhost),
          'Server' => true,
          'Context' =>
            {
              'Msf' => framework,
              'MsfExploit' => self
            }
        )

        unless opts[:logger]
          log_device = LogAdapter::LogDevice::Framework.new(framework)
          opts[:logger] = LogAdapter::Logger.new(self, log_device)
        end

        thread_factory = Proc.new do |server_client, &block|
          Rex::ThreadFactory.spawn("SMBServerClient(#{server_client.peerhost}->#{server_client.dispatcher.tcp_socket.localhost})", false, &block)
        end

        server = RubySMB::Server.new(
          server_sock: @rsock,
          gss_provider: opts[:gss_provider],
          logger: opts[:logger],
          thread_factory: thread_factory
        )

        server.extend(ServiceMixin)
        server.on_client_connect_proc = Proc.new { |client|
          on_client_connect(client)
        }
        self.service = server
        self.service.start

        print_status("Server is running. Listening on #{bindhost}:#{bindport}")
      end

      def on_client_connect(client)
        vprint_status("Received SMB connection from #{client.peerhost}")
      end

      def cleanup_service
        if service
          begin
            self.service.stop
            self.service.wait
            true
          rescue ::Exception => e
            print_error(e.message)
            false
          end
        end
      end
    end
  end
end

